EmbedID に学習済み word2vec の重みパラメータをセットする [Chainer]

深層学習における自然言語処理のタスクでは、単語の埋め込み層として、学習済み word2vec (Skip-gram / CBOW) モデルを使用することがあります。 本記事では、Chainer で EmbedID に学習済み word2vec をセットする方法を説明します。

例として、学習済みモデルは次のページで配布されているものを使用します。
GitHub - Kyubyong/wordvectors: Pre-trained word vectors of 30+ languages

上記のページにある "Japanese (w)" のリンクから word2vec モデルをダウンロードして解凍します(Requirements に "mecab" とあるので、単語分割には MeCab + IPA 辞書を使用しているのでしょうか)。

次のコードで EmbedID に学習済み word2vec の重みパラメータをセットします。
※ あらかじめ、Chainer, gensim, NumPy をインストールしておいてください

import gensim
import numpy as np
import chainer.links as L
from chainer import Variable

# word2vec モデルを読み込む
word2vec = gensim.models.Word2Vec.load('ja/ja.bin')

# 単語辞書 (index to word / word to index) を用意する
i2w = word2vec.wv.index2word
w2i = {w: i for i, w in enumerate(i2w)}
n_vocab = len(i2w)

# EmbedID を初期化し、重みパラメータとして word2vec をセットする
emb = L.EmbedID(n_vocab, word2vec.vector_size, initialW=word2vec.wv.vectors)

emb.disable_update()

最終行の emb.disable_update() を実行すると、ニューラルネットの学習中に EmbedID の重みが更新されなくなります。

これで準備は完了です。確認のために、gensim と Chainer で、それぞれ同じ単語を与えたときに同じベクトルが取得できるか見てみます。

e1 = word2vec['今日']
print(e1)

# [-1.5841931  -0.16819568  0.17484002 -0.70384747 -0.34765592 -0.355766
#  -0.23339556  0.6226926   0.2746878   0.15454574 -0.70101714 -0.5748118
#  ...
#  -0.42160854 -0.44855088  1.7454495  -1.893566    1.4198955   0.12920022
#   2.1218472  -0.7616498   3.1082888  -1.3034075   0.2861628   0.31217498]

x = Variable(np.array([w2i['今日']], dtype=np.int32))
e2 = emb(x)
print(e2[0].data)

# [-1.5841931  -0.16819568  0.17484002 -0.70384747 -0.34765592 -0.355766
#  -0.23339556  0.6226926   0.2746878   0.15454574 -0.70101714 -0.5748118
#  ...
#  -0.42160854 -0.44855088  1.7454495  -1.893566    1.4198955   0.12920022
#   2.1218472  -0.7616498   3.1082888  -1.3034075   0.2861628   0.31217498]

同じベクトルが取得できました。問題なく動作するようです。

バッチでベクトルの内積を計算する [Chainer]

Chainer(本記事の執筆時点で最新版は v4.3.0)の標準関数でサポートされていない(と思う)のでメモ。

次のコードでベクトルの内積が計算できます。

import numpy as np
import chainer.functions as F
from chainer import Variable

def inner_product(a, b):
    return F.sum(a * b, axis=1)


# [0.1, 0.2] と [-0.5, 0.4] の内積 -> 0.03
# [0.3, 0.4] と [-0.3, 0.2] の内積 -> -0.01
a = Variable(np.array([[0.1, 0.2], [0.3, 0.4]], dtype=np.float32))
b = Variable(np.array([[-0.5, 0.4], [-0.3, 0.2]], dtype=np.float32))

inner_product(a, b)
# => variable([ 0.03 -0.01])

もうこれで十分な気がしますが、一応詳細を説明します。

ベクトル  \boldsymbol{a} = (a_1, a_2, \cdots, a_n)^\top \boldsymbol{b} = (b_1, b_2, \cdots, b_n)^\top内積は次のように表されます。

{\displaystyle \boldsymbol{a} \cdot \boldsymbol{b} = \sum_{i=1}^n a_i b_i}

つまり、各成分の積を足していけば OK です。

冒頭のコードでは、a * b で要素(成分)ごとの積を求めています。そして、F.sum()axis=1 として、バッチ内のインスタンス単位で要素の総和を計算しています。

余談: 要素積を求めるのに * 演算子を使えば良いことに気づくまでしばらくかかりました。該当する関数が無いか、かなりドキュメントを漁っていました…。

NStepLSTM の前段で EmbedID を使う [Chainer]

NStepLSTM を使用したニューラルネットを構築する場合、NStepLSTM に対して単語の埋め込みベクトルを入力したい場合があると思います。

NStepLSTM の入力 (xs) のデータ型は Variable のリストであり、これは、テキスト系のタスクにおいて、文の長さが可変であることを踏まえた構造になっているといえます。

しかし、EmbedID の(入力および)出力のデータ型は Variable であり、NStepLSTM と扱うデータ型が少々異なるので、これら2つを組み合わせて使用したい場合は一工夫が必要です。

次のコードはエラーの例です。

import numpy as np
import chainer.links as L
from chainer import Variable

# 適当なニューラルネット (注意: 本来はこのような小規模なネットワークは組みません)
embed = L.EmbedID(10, 3)
lstm = L.NStepLSTM(2, 3, 3, 0.5)

# 各文の単語 ID 列を Variable でラップする
xs = [
  Variable(np.array([0, 5, 7, 1], dtype=np.int32)),
  Variable(np.array([0, 4, 3, 6, 7, 1], dtype=np.int32))
]

# ここでエラーが発生
emb_xs = embed(xs)

hy, cy, ys = lstm(None, None, emb_xs)

これを解決するために、次の関数を導入します。 ※この関数は公式 Chainer の Example (text_classification/nets.py, seq2seq/seq2seq.py) で使用されているものです。

import chainer.functions as F

def sequence_embed(embed, xs):
    x_len = [len(x) for x in xs]
    x_section = np.cumsum(x_len[:-1])
    ex = embed(F.concat(xs, axis=0))
    exs = F.split_axis(ex, x_section, 0)
    return exs

この関数と EmbedID を組み合わせれば OK です。sequence_embed() の第1引数に EmbedID のオブジェクト、第2引数に EmbedID への入力を Variable のリストの形式で渡します。例えば、次のようなコードになります。

embed = L.EmbedID(10, 3)
lstm = L.NStepLSTM(2, 3, 3, 0.5)

# 各文の単語 ID 列を Variable でラップする
xs = [
  Variable(np.array([0, 5, 7, 1], dtype=np.int32)),
  Variable(np.array([0, 4, 3, 6, 7, 1], dtype=np.int32))
]

# 正常に処理される
emb_xs = sequence_embed(embed, xs)

hy, cy, ys = lstm(None, None, emb_xs)

sequence_embed() 内部の処理をざっくりと説明すると、Variable のリストを Flat な Variable に変換 → EmdedID に通す → その結果を再び Variable のリストに戻す、という感じのことをやっています。

本記事では NStepLSTM を取り扱いましたが、NStepBiLSTM, NStepGRU, NStepBiGRU についても同様です。

サンプルコード完全版は こちら