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 についても同様です。

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