kerasで学習済みword2vecをモデルに組み込む方法を紹介します。word2vecなどで学習した分散表現（token id毎のベクトル値）をkerasのembedding layerの重みに設定し、新たに学習させないように指定するという流れです。こうすることで、word2vecによる特徴量抽出を行うモデルがkerasで構成できます。自然言語処理では頻出の処理だと思います。前処理の段階で分散表現を獲得する(layerに組み込まない)手もあるかとは思いますが、prediction用に別途前処理用のコードを書く必要がでてくるなどの手間が発生するので、できるだけ今回紹介するようにモデルに加えるほうが無難だと思います。

使用するフレームワーク

分散表現獲得: gensim v3.8.0
NN構築: keras(tensorflow backend) v2.2.5

データセット

IMDB Movie reviews sentiment classification datasetを使用します。映画のレビューコメントの感情分析(positive or negative)のアノテーションデータです。keras.datasetsから簡単に取得できます. Datasets （今回はモデルには学習させないのでデータセットは何でもいいです）

取得するデータは前処理済みのtoken id(int)のリストになっています。gensimの仕様に合わせるためにここでtoken idを文字列に変換しておきます。

from keras.datasets import imdb

num_words = 1000
(x_train, y_train), (x_test, y_test) = imdb.load_data(path="imdb.npz",
                                                      num_words=num_words)
x_train = [[str(f) for f in x] for x in x_train] # ID文字列のリストに変換

word2vecでの分散表現の学習

gensimを使ってword2vecの学習を行います。wikipediaデータなどの大規模コーパスを使用するのが一般的だと思いますが、ここでは簡単のためにtrain dataを使います。

from gensim.models import Word2Vec

 # 分散表現の次元
w2v_size = 100
 # 学習開始
w2v = Word2Vec(x_train, size=w2v_size, min_count=3, window=3)

gensimのword2vecからembedding layer用の重み行列を作る

先程学習したモデルはデータセットの全てのtoken idに対する分散表現(ベクトル値)をもっています。これをembedding layerの重みにします。すると、モデルにword2vecを組み込んだ形になります。ただし、keras.layers.Embeddingで指定するために行列の形式に変換する必要があります。

def get_embedding_matrix(model, word_index):
    """
    keras.layers.Embeddingのweights引数で指定するための重み行列作成
    model: gensim model
    num_word: modelのvocabularyに登録されている単語数
    emb_dim: 分散表現の次元
    word_index: gensim modelのvocabularyに登録されている単語名をkeyとし、token idをvalueとする辞書 ex) {'word A in gensim model vocab': integer token id} 
    """
    # gensim modelの分散表現を格納するための変数を宣言
    embedding_matrix = np.zeros((max(list(word_index.values())) + 1, model.vector_size), dtype="float32")
   
    # 分散表現を順に行列に格納する
    for word, label in word_index.items():
        try:
            # gensimのvocabularyに登録している文字列をembedding layerに入力するone-hot vectorのインデックスに変換して、該当する重み行列の要素に分散表現を代入
            embedding_matrix[label] = model.wv[word]
        except KeyError:
            pass
    return embedding_matrix

実行例

上記関数は以下のように呼び出します。

# word_indexの作成
word_index = {key: int(key) for key in w2v.wv.vocab}

# ゼロ埋め用(入力の時系列方向の長さを揃えるためにtoken id=0でpaddingする想定)にindexを追加
word_index['0'] = 0

# embedding_matrixを作成
emb_mat = get_embedding_matrix(w2v, word_index)

emb_mat.shape
# >>> (1000, 100)

keras.layers.Embeddingの宣言法

先程作成したembedding_matrixをweighs引数に代入します。また、embedding layerが追加で学習しないようにtrainable=Falseとします。

from keras.layers import Embedding

# 単語数、embeddingの次元
num_words, w2v_size = embedding_matrix.shape

# embedding
Embedding(num_words, w2v_size,
                        weights=[embedding_matrix], 
                        input_length=max_len, # 時系列方向の大きさ(最大入力単語数)
                       trainable=False)

modelの例

LSTMを使ったシンプルなモデル例を挙げておきます。

# 最大入力単語数
max_len = pd.DataFrame(x_train)[0].apply(lambda x: len(x)).max()
# >> 2494

# 単語数、embeddingの次元
num_words, w2v_size = embedding_matrix.shape

# モデルの構築
model = Sequential()
model.add(Embedding(num_words, w2v_size,
                    weights=[embedding_matrix], 
                    input_length=max_len,
                   trainable=False))
model.add(LSTM(50, batch_input_shape=(None, max_len, w2v_size), return_sequences=False))
model.add(Dense(25))
model.add(Dense(2))
model.add(Activation("sigmoid"))

model.summary()

_________________________________________________________________
Layer (type)                 Output Shape              Param #   
=================================================================
embedding_2 (Embedding)      (None, 2494, 100)         100000    
_________________________________________________________________
lstm_1 (LSTM)                (None, 50)                30200     
_________________________________________________________________
dense_1 (Dense)              (None, 25)                1275      
_________________________________________________________________
dense_2 (Dense)              (None, 2)                 52        
_________________________________________________________________
activation_1 (Activation)    (None, 2)                 0         
=================================================================
Total params: 131,527
Trainable params: 31,527
Non-trainable params: 100,000
_________________________________________________________________

サマリーを見るとembeddingのParam #=100000であり、Non-trainable params: 100000なのでembedding layerのフリーズ（重みを学習しない）は成功していそうです。gensimを使った分散表現を用いると、大幅にNeural Networkの学習を行う重みの数を削減できることが確認できました。 BERT以前はword2vecで学習済みのembedding layerを導入することで大幅に精度が向上したそうです。既に一般的な手法であり、ベースラインとして使ったり、リソースをあまり割きたくない時に便利な手法ですので、ぜひ興味があればお試し下さい！