Tensorflow2.1で追加されたTextVectorizationで文字列から直接訓練できるようになる

この記事はfreeeデータに関わる人たち Advent Calendar 2019の15日目のエントリーです。

データ分析とは関係ない小ネタをお送りしてまいります。

もうそろそろTensorflow v2.1がリリースされます。v2.1ではtensorflowとtensorflow-gpuが一つのパッケージに統合されたり（今後は軽量なCPU only versionとの二本立てになる模様）、様々なバグが修正されているのでお使いの皆さんは早々に乗り換えるのがいいと思います。

本日は、実験的にサポートされたkerasのLayerであるTextVectorizationについて紹介します。

今までTensorflowで文字列を入力に使用する場合には、文字列を前処理で数値列に変換しておく必要がありました。例えばこんな感じです：

# 訓練時

dataset = load_dataset("dataset_path")  # データを読み込む
tokenized_text = tokenize_text(dataset.text_feature)  # 文字列をToken（単語）の塊にする
token_id_map = build_id_map(tokenized_text)  # Tokenを数値に変換するマップを作る
save_metadata(token_id_map, "metadata_file_path")  # 保存しておいて推論時などに使う

token_ids = token_to_ids(tokenized_text, token_id_map)
train(token_ids)  # 学習に使う

# 推論時

model = load_model("model_path")
token_id_map = load_metadata("metadata_file_path")
input_data = token_to_ids(tokenized_text, token_id_map)
model.predict(input_data)

文字列をTokenizeした後、それをテーブルを使って数値IDに変換するのが大変でした。

今までも一応StaticVocabularyTableというものを使えば変換テーブル自体はモデルに組み込んだりできたのですが、tf.kerasで使うとなると自分でカスタムLayerを作る必要はあるし、地味に面倒くさかったのです。

TextVectorizationは若干の制約はありますが、TokenizationからID変換までを一気にやってくれるという、ものぐさな人にピッタリの便利Layerになっています。

使い方

公式にもサンプルがあるのでそちらもご覧になるといいかと思います。

1. TextVectorationにテキストを入力して、マッピングをあらかじめ作成する

import tensorflow as tf
from tensorflow.keras import layers
from tensorflow.keras.layers.experimental.preprocessing import TextVectorization
import numpy as np

train_data = np.array(['おはよう', 'こんにちは', 'さようなら'])
label = np.array([0, 0, 1])

def _normalize(input_data):
    """文字列のカスタム正規化処理ができる"""
    lowercase = tf.strings.lower(input_data)
    return lowercase

def _split(input_data):
    """文字列のTokenizeができる。やりたい人はtf-sentencepieceとかを頑張って使えるかも（未検証）"""
    # 今回は文字単位でToken化する
    return tf.strings.unicode_split(input_data, input_encoding='UTF-8')

text_vectorize_layer = TextVectorization(
    max_tokens=1000, 
    standardize=_normalize, 
    split=_split, 
    ngrams=(1, 2), 
    output_mode='int', 
    output_sequence_length=40
)
text_vectorize_layer.adapt(train_data)

もちろん大きいデータセットの場合に備えてストリームで読み込みも可能です。tf.dataを使ってもいいですし、自分で実装してもいいでしょう。

TextVectorizationのオプションは意味は以下のようになってます（ほかにもある）

<max_tokens>
Vocaburaryの上限を指定（Noneなら無制限）

<standardize & split>
文字列の正規化処理とTokenize処理をカスタム可能

<ngrams>
tokenのngramをそのままID可出来る（指定したgram順にベクトル化されて出力される）

<output_mode>
int, binary, count, tf-idfから選べる。intはテキストを ID Sequenceで出力する。binary, count, tf-idfは、サイズがVocabularyのベクトルを出力し、binaryがone-hot表現（出現したら1）、countはTokenの出現数、tf-idfはTokenのtf-idfが対応するインデックスの値に入る。

<output_sequence_length>
ouotput_mode="int"の時のみ必要。出力するベクトルの長さ=最大受付Token数。この値をはみ出した分は多分無視（後ろ側がtruncate）されると思うが未確認

N-gramが最初から使えたり、one-hotかsequenceか選べたり、tf-idfが使えたりと結構便利です。

2. tf.kerasのlayerとして使う

input_layer = layers.Input(shape=(1,), dtype=tf.string, name='input')
text_vector = text_vectorize_layer(input_layer)
embed = layers.Embedding(input_dim=1001, output_dim=100, input_length=40)(text_vector)
output_layer = layers.Dense(1, activation='tanh')(embed)
model = tf.keras.Model(inputs=input_layer, outputs=output_layer)

サンプルなのでとても単純です。ここでのポイントは、

output_mode="int"を指定しており、続けてlayers.Embeddingでembedding vectorに変換していること

そのパラメータはinput_dim=max_tokens+1（1-originでID化されるので必ず+1する必要あり）、input_length=output_sequence_lengthにする

の二点です。

3. 訓練する

model.compile(loss='binary_crossentropy', optimizer='adam', metrics=['accuracy'])
model.fit(x=train_data, y=label, epochs=5)

train_dataは文字列の配列ですが、これでちゃんと訓練できます。とてもコードがスッキリしていい感じです。

以上です。TextVectorization Layerを使えば、簡単なモデル作りであれば爆速で実装できそうな気配がします。しかもこのモデル、save/loadが簡単にできるので、辞書ファイルを別途読み込んだりしなくて済むのが地味にうれしいです。未検証ですが多分TF servingでも使えると思います。