TF2.0 tf.keras.Model 用 CRFレイヤーを試してみた

はじめに

こんにちは。 リサーチャーの坂田です。 レトリバでは、 主に文書分類とNERに取り組んでいます。

tensorflow 2.0 では、 1系のときに"tf.contrib"以下にあったcrfが、 "tensorflow_addons"として分離されました。 現在、 tf.kerasのレイヤーとして使えるように開発が進められています。 今回は、 このcrfの最新のプルリクを試してみたという内容です。まだマージされていないものなので、色々と問題がありました。

CoNLL 2003というデータセットを使って、 NERのタスクでCRFを試しました.

NER とは

NER (Named Entity Recognition)とは、 文書中から、 あらかじめ定められたカテゴリに当てはまる固有表現を抽出するタスクです。 1番有名なデータセットとしては、 CoNLL 2003があります。CoNLL 2003では、 文中の固有表現に対して、 persons(人物)、organizations(組織)、 locations(場所)、 miscellaneous(その他)のラベルを付けます。 NERの研究では、 よく評価データセットとして使われています。英語とドイツ語が用意されています。

例えば、以下のようなラベル付けが考えられます。

Taro Yamada works as an engineer for Retrieva .
B-PER I-PER O O O O O B-ORG O

このように、固有表現の始端トークンにB-(固有表現の種類)、 それ以降の固有表現内トークンにI-(固有表現の種類)、固有表現以外のトークンにはOのラベルが付与されます。上記の例ではPERは人物を、ORGは組織を表します。 これには、いくつかのラベル付けがあり、上記は一番簡単なラベル付けです。

コールセンター向けシステムでは、 会話データに対して、 人物名、 会社名、 日時、 製品名などのラベルを付与することが考えられます。

CRF とは

CRF (Conditional Random Field) とは、 系列ラベリング系タスクで良く使われる手法で、系列の前後関係を考慮したラベル付けを行います。 NERでは、 単語やラベルの前後関係が大きく影響すると思われるため、 CRFが使われることが多いです。CRFについて詳しく知りたい方は、以下の資料や書籍などが分かりやすいと思います。

tensorflow 2.0

Tensorflowは、 Googleが2015年にリリースしたディープラーニングフレームワークです。 1系時代は、 慣れるまでがとても大変と言われていましたが、 2系になってインタフェースが刷新されました。 Chainer で採用された Define by Run を取り入れており、 大幅に使いやすくなっています。 先日、 Chainerの開発が停止されるとの発表がありました. PyTorchは、 初期はChainerをフォークして開発されていましたし、 Tensorflowの2系もその影響を強く受けています。

tf.keras

tf.kerasは、 tensorlfowで簡単にモデルを構築するための高レベルAPIです。tensorflow 1系時代は、 tf.estimatorが使われることが多かった印象ですが、 2系に移行するにあたり、 Googleは tf.kerasの方を推しているようです。Kerasは、 元々theanoというディープラーニングフレームワークのラッパーとして登場し、 後にtensorflowにも対応しました. 現在では、 tf.kerasとしてtensorflowに取り込まれています。keras と tf.kerasで挙動が異なる場合があるため、 ネット上で調べるときには注意が必要です。 tf.kerasにはいくつか使い方があります。 - tf.keras.Sequentialを使う方法 - Functional API を使う方法 - tf.keras.Modelを継承する形でモデルを定義する方法

いずれの場合でも、 モデルの作成後以下の1.~4.の手順で学習、 評価を行います。 ただし、 独自に作成したレイヤーを用いた場合や、 tf.keras.Modelを継承してモデルを定義した場合に、 この手順が使えないことがあります。そのときは、 自分で訓練ループを書く必要があります。 (詳しくは公式ドキュメント参照) 1. compileメソッドでモデルの初期化 2. summaryメソッドで確認 (省略可) 3. fitメソッドで訓練 4. evaluateメソッドで評価

tf.keras.Sequentialを使う方法

一番簡単なのは、 以下のように tf.keras.Sequentialを使ってモデルを定義する方法です。 まず、 tf.keras.Sequentialオブジェクトを作成し、 そこにaddメソッドを使ってレイヤーを追加していきます。 リストにappendするのと同じ感覚で簡単にモデルを作成できます。 (公式ドキュメント)

import tensorflow as tf 

model = tf.keras.Sequential()
model.add(tf.keras.layers.Dense(32, activation='relu'))
model.add(tf.keras.layers.Dense(10, activation='softmax'))

Functional API を使う方法

テキストと画像など、 ネットワークへの入力が複数の場合や、 エンコーダーデコーダーモデルで、 エンコーダーデコーダーの重みを共有する必要があるといった場合に便利です。 NERで、 文字の分散表現と単語の分散表現を入力とするといった場合が、 これに該当します。まずtf.keras.Inputで入力レイヤーを定義し、 その後、 各レイヤーのcallメソッドに前のレイヤーの出力を渡す形でネットワークを定義します。 (インスタンス名に()を付けるとcallメソッドが呼ばれる。) (公式ドキュメント)

import tensorflow as tf 

input_1 = tf.keras.Input(shape=(None,))
input_2 = tf.keras.Input(shape=(None,))

embedding_1 = tf.keras.layers.Embedding()(input_1)
embedding_2 = tf.keras.layers.Embedding()(input_2)

embedding = tf.keras.layers.concatenate([embedding_1, embedding_2])

x = tf.keras.layers.Dense(32, activation='relu')(embedding)
x = tf.keras.layers.Dense(10, activation='softmax')(x)

tf.keras.Modelを継承する形でモデルを定義する方法

さらに自由度の高いモデルを作りたいという場合は、 以下のようにtf.keras.Modelを継承したクラスを作成します。この方法ならば、 インスタンス変数として独自の重みを保持できるため、 さらに自由度が上がります。 (公式ドキュメント)

tensorflow_addons

tensorflow 2.0では、 コアに含まれない追加の機能がtensorflow_addonsとして提供されています。 CRFレイヤーもこの中に含まれています。インストールは pip で可能です。

pip install tensorflow_addons

実装

以下、 NERの実装です。 コードはgithubに置いています。

使用したCoNLL 2003 データセットは、 こちらのgithubリポジトリからダウンロードしました。

単語ベクトルは、 gensimでダウンロード出来るpretrainモデル(glove-wiki-gigaword-300)を使いました。 詳しくは、 githubを御覧ください。

また、 以下のようにプルリクをダウンロードしました。

git clone https://github.com/tensorflow/addons.git
git fetch origin pull/377/head:crf_layer 
git checkout crf_layer

setup.pyはエラーが出て動かなかったため、 text、 layers、 lossesの3つのディレクトリをpipで入れたtensorflow_addonsのディレクトリ以下に上書きすることでインストールしました。

動作環境 - docker image tensorflow/tensorflow:2.0.0-gpu-py3

import datetime

import numpy as np
import tensorflow as tf 
from tensorflow.keras import Sequential
from tensorflow.keras.layers import Bidirectional, Dense, Embedding, LSTM
from tensorflow.keras.preprocessing.sequence import pad_sequences
from tensorflow_addons.layers import CRF 
from tensorflow_addons.losses import ConditionalRandomFieldLoss
import tensorflow_datasets as tfds

from conll_dataset import ConllLoader

BATCH_SIZE = 5 # 訓練時のバッチサイズ
EPOCHS = 5 # 訓練時のエポック数
EMBEDDING_DIM = 300 # 学習済み単語分散表現の次元数
MAX_LENGTH = 32 # 学習データの最大系列長より大きい必要がある

# CoNLL 2003 データセットから(stringのlist, stringのlist)の形式で、各種データを読み込みます。
X_train, y_train = ConllLoader('/mnt/NER/corpus/CoNLL-2003/eng.train').get_list()
X_dev, y_dev = ConllLoader('/mnt/NER/corpus/CoNLL-2003/eng.testa').get_list()
X_test_original, y_test = ConllLoader('/mnt/NER/corpus/CoNLL-2003/eng.testb').get_list()

# 以下で学習済み分散表現の読み込みを行います。
# これによって、訓練データ以外の大規模コーパスから得た単語の意味情報を考慮できます。
# 0番目はパディング、 (語彙数+1)番目は未知語に割り当てられます。
glove = np.concatenate([np.zeros(EMBEDDING_DIM)[np.newaxis], np.load('../glove-wiki-300-connl.npy'), 
        np.zeros(EMBEDDING_DIM)[np.newaxis]], axis=0)

# 学習済み単語分散表現の語彙リストを読み込みます。
with open('../glove-wiki-300-connl.vocab') as f:
    vocab_list = f.readlines()

# スペース区切り単語列を単語インデックス列に変換するエンコーダを定義。
word_token_encoder = tfds.features.text.TokenTextEncoder(vocab_list)

# 全ラベルリスト と ラベル数 の取得
label_list = list(set([label for label_in_sentence in y_train for label in label_in_sentence.split()]))
n_labels = len(label_list)

# スペース区切りラベル列を単語インデックス列に変換するエンコーダを定義。
label_token_encoder = tfds.features.text.TokenTextEncoder(label_list)

def encode_and_pad_data(sequence_list, encoder):
    # strのlistを単語インデックスのlistのlistに変換した後、
    # MAX_LENGTHに満たない文に対して、0でパディングを行う。
    enceded_sequences = [encoder.encode(sequence) for sequence in sequence_list]
    encoded_and_padded_sequences = pad_sequences(enceded_sequences, maxlen=MAX_LENGTH, dtype='int32', padding='post', value=0)
    return encoded_and_padded_sequences


# 以下でエンコードとパディングを行います。
X_train = encode_and_pad_data(X_train, word_token_encoder)
X_dev = encode_and_pad_data(X_dev, word_token_encoder)
X_test = encode_and_pad_data(X_test_original, word_token_encoder)
y_train = encode_and_pad_data(y_train, label_token_encoder)
y_dev = encode_and_pad_data(y_dev, label_token_encoder)
y_test = encode_and_pad_data(y_test, label_token_encoder)


# 以下でモデルの定義
# Sequential APIを使ってモデルの定義
model = tf.keras.Sequential()
# Embedding レイヤーで、 単語インデックスをそれに対応する単語分散表現に置き換えます。
model.add(Embedding(input_dim=word_token_encoder.vocab_size, output_dim=EMBEDDING_DIM, 
                    embeddings_initializer=tf.keras.initializers.Constant(glove), 
                    input_shape=(MAX_LENGTH, )))# mask_zero=True は上手くいかない
# 出力次元数200のBiLSTMレイヤーを追加します。
# LSTMクラスのイニシャライザにreturn_sequences=Trueを渡すことで、
# 全ての単語に対する予測値を取り出せます。
# return_sequences=Falseだと1文に対して1つの予測値のみの出力となります。
model.add(Bidirectional(LSTM(200, return_sequences=True), merge_mode='ave'))
# 出力次元数がタグの種類数の全結合層を追加します。
# また、活性化関数をsoftmaxにします。
# これは、CRFレイヤーが各ラベルに対する予測の確信度を要求するためです。
model.add(Dense(n_labels, activation='softmax'))
# 最後にCRFレイヤーを追加します。
model.add(CRF(n_labels, name='crf_layer'))

# モデルの初期化を行います。ここで、誤差関数にConditionalRandomFieldLossを指定します。
model.compile(optimizer='adam', loss=ConditionalRandomFieldLoss())
# TensorBoard用のログの保存場所を定義します。
log_dir="logs/fit/" + datetime.datetime.now().strftime("%Y%m%d-%H%M%S")
# ここで、学習データを渡すことで、学習を行います。
model.fit(X_train, y_train, validation_data=(X_dev, y_dev), epochs=EPOCHS,  
    batch_size=BATCH_SIZE, 
    callbacks=[tf.keras.callbacks.TensorBoard(log_dir=log_dir, histogram_freq=1)])
# モデルの保存
model.save('saved_model/model_sequential.h5')

# 予測ラベルの取得
y_pred = model.predict([X_test][:5])

# 予測結果の表示
for i in range(5):
    print('元の文:', X_test_original[i])    
    print('正解ラベル:', label_token_encoder.decode(y_test[i]))    
    print('予測ラベル:', label_token_encoder.decode(y_pred[i]))  
    print()  
    with open('ner_results.txt', 'a') as f:
        f.write('元の文:'+X_test_original[i]+'\n')
        f.write('正解ラベル:'+label_token_encoder.decode(y_test[i])+'\n')
        f.write('予測ラベル:'+label_token_encoder.decode(y_pred[i])+'\n\n')

ラベル予測結果

元の文                 SOCCER - JAPAN GET LUCKY WIN , CHINA IN SURPRISE DEFEAT .
正解ラベル O O I-LOC O O O O I-PER O O O O
予測ラベル O O O O O O O O O O O O
元の文 Nadim Ladki
正解ラベル I-PER I-PER
予測ラベル O O
元の文                 AL-AIN               , United Arab Emirates 1996/12/6
正解ラベル I-LOC O I-LOC I-LOC I-LOC O
予測ラベル O O I-ORG O O I-ORG O O O O O
元の文                 Japan began the defence of their Asian Cup title with a licky 2-1 win against Syria in a Group C championship match on Friday .
正解ラベル I-LOC O O O O O I-MISC I-MISC O O O O O O O I-LOC O O O O O O O O O
予測ラベル O O O O O O I-MISC I-MISC O O O O O O O I-MISC O O O O O O O O O
元の文                 But China saw their luck desert them in the second match of the group , crashing to a surprise 2-0 defeat to newcomers Uzbekistan .
正解ラベル O I-LOC O O O O O O O O O O O O O O O O O O O O O I-LOC O
予測ラベル O I-MISC O O O O O O O O O O O O O O O O O O O O O O O
accuracy:  16.43%; precision:   7.05%; recall:  10.41%; FB1:   8.41
              LOC: precision:   0.23%; recall:   0.78%; FB1:   0.36  5624
             MISC: precision:  12.51%; recall:  15.53%; FB1:  13.86  871
              ORG: precision:  25.71%; recall:  27.87%; FB1:  26.75  1801
              PER: precision:   6.82%; recall:   0.19%; FB1:   0.36  44

結果の出力にはこちらのgithubリポジトリの評価スクリプトを使わせて頂きました。model_sequential_evaluation.pyを使って、評価スクリプトに合わせた形式に結果を変換出来ます。

一部当たっているところもありますが、パディングをマスクして無視する機能が使えないこともあり、 全体のF値のミクロ平均が10.0%を切ってしまいました。

現状で起こる問題

tensorflow addons における CRFの現在の実装では、 Embeddingレイヤーのmask_zeroオプションを適用できません。 以下のようなエラーが出ます。

TypeError: Value passed to parameter 'x' has DataType bool not in list of allowed values: float32, float64, int32, uint8, int16, int8, int64, bfloat16, uint16, float16, uint32, uint64

これは、 CRFレイヤーには、 前層から[True, True, True, ... , False, False]といったマスクキングが渡され、 レイヤー内の演算がbooleanに対応していないことが原因です。 tensorflow_addonsのissueでもこのことが指摘されており、 CRFレイヤーの内部で型をキャストする方法が示されていますが、 以下のようなtensorflowのコア部分のエラーが出て、 解決に至りませんでした。

tensorflow_addons/layers/crf.py:243 call  *
        left_boundary_mask = self._compute_mask_left_boundary(mask)
    tensorflow_addons/layers/crf.py:335 _compute_mask_left_boundary  *
        left_boundary = tf.keras.backend.greater(tf.dtypes.cast(mask, tf.int32), right_shifted_mask)
    tensorflow_core/python/keras/backend.py:2371 greater
        return math_ops.greater(x, y)
    tensorflow_core/python/ops/gen_math_ops.py:4398 greater
        "Greater", x=x, y=y, name=name)
    tensorflow_core/python/framework/op_def_library.py:563 _apply_op_helper
        inferred_from[input_arg.type_attr]))

    TypeError: Input 'y' of 'Greater' Op has type bool that does not match type int32 of argument 'x'.

また、 tf.keras.Modelを継承する形でモデルを定義する方法で、 CRFレイヤーを含むモデルを定義しようとすると、 さらなる問題があります。 CRFは、 lossを計算するときに、 予測ラベルと正解ラベルだけではなく、 遷移行列を必要とします。 それため、 tensorflow_addons.losses.ConditionalRandomFieldLossでは、 予測値のKeras Tensorオブジェクト(tf.Tensorとは別) の _keras_historyプロパティからCRFレイヤーにアクセスしますが、 これが上手くいかない場合があります。CRFレイヤーの開発者によって回避法は示されてはいますが、 モデルの中で分岐がある場合(複数入力や、 concatを含む、 etc)は別途書き換えが必要なようです。 より複雑なモデルを実装しようとするとこちらの手法が必要になりますが、 Tensorflow のコア部分の実装が影響しているようで、 解決には時間が掛かるかもしれません。

以下のように、 モデルに、 loss関数を定義してしまう方法も試してみましたが、 この方法も上手くいきません。

# tf.keras.Model を継承してモデルを定義する。
class BilstmCRF(tf.keras.Model):
    def __init__(self, num_tags, token_encoder, label_encoder, embedding_list):
        super().__init__()
        self.num_tags = num_tags
        self.token_encoder = token_encoder
        self.label_encoder = label_encoder

        self.embedding_layer = Embedding(input_dim=token_encoder.vocab_size, output_dim=EMBEDDING_DIM, 
                                         embeddings_initializer=tf.keras.initializers.Constant(embedding_list), mask_zero=True)
        self.bilstm_layer = Bidirectional(LSTM(200, return_sequences=True), merge_mode='ave')
        self.dense_layer = Dense(num_tags, activation='softmax')
        self.crf_layer = CRF(num_tags, name='crf_layer')

    @tf.function(autograph=True) # autograph=Falseでグラフモードへの変換を停止
    def call(self, inputs, training):
        x = self.embedding_layer(inputs)
        mask = x._keras_mask
        x = self.bilstm_layer(x, mask=mask, training=training)
        x = self.dense_layer(x, training=training)
        x = self.crf_layer(x, mask=tf.cast(mask, tf.int32), training=training) 

        return x

    def crf_loss(self, true_label, pred_label):
        loss = self.crf_layer.get_loss(true_label, pred_label)
        return loss

callメソッドを、 tf.functionでデコレートすることで、 グラフモード/define and run形式(tf1系と同じ)で動作します。 これによって、 動作速度の向上が見込まれます。 しかし、 この場合は、 CRFレイヤー内のget_lossメソッドにアクセス出来なくなり、 以下のエラーがでます。

tensorflow.python.framework.errors_impl.InaccessibleTensorError: The tensor 'Tensor("crf_layer/Cast_3:0", shape=(None,), dtype=int64)' cannot be accessed here: it is defined in another function or code block. Use return values, explicit Python locals or TensorFlow collections to access it. Defined in: FuncGraph(name=call, id=140226998764752); accessed from: FuncGraph(name=keras_graph, id=140229058436624).

また、 tf.functionを無効にした場合でも、 以下のようなエラーが出て上手くいきません。

ValueError: Can not squeeze dim[1], expected a dimension of 1, got 32 for 'metrics/accuracy/Squeeze' (op: 'Squeeze') with input shapes: [?,32].

おわりに

Tensorflow 2.0で、 tensorflow_addonsのCRFレイヤーのプルリクを試しました。 まだ問題があってマージされていないのだから当然ではありますが、 上手く動作しないという結果となりました。 今回の取り組みの中で、 tensorflowのコア部分のコードを読むことで大変勉強になりました。

Tensorflow 2.0 自体は、 1系と比べて とても使いやすくなっていますが、 CRFを導入するのは なかなか大変なようです。 現状でCRFを使いたい場合は、モデルの最終層をDenseレイヤーにして、TF1系と同じようにtensorflow_addons.text.crf_log_likelihoodを使用する方法があります。

参考文献