自然言語処理の深遠

Deep Dive Into Natural Language Processing

Keras の RNN/LSTM/GRU で内部状態を取得する

自然言語処理で RNN を使っていると、RNN の内部状態を取得したくなることがあります。 TensorFlow では tf.nn.dynamic_rnn 等の関数を使うと、出力と状態を返してくれます。 しかし、Keras でのやり方については意外と日本語の情報がありませんでした。

本記事では Keras で RNN の内部状態を取得する方法についてまとめてみました。

RNN/LSTM/GRU の内部状態を取得

Keras にはリカレント層として、SimpleRNNLSTMGRU の3種類が用意されています。これらの層から内部状態を取得するためには、インスタンス化時の引数として return_state=True を渡す必要があります。 return_state を True にすることで、出力に加えて最終状態を取得できるようになります。

では、実際に各リカレント層の内部状態を取得してみましょう。 まずは各層への入力を作ります。 ここでは簡単に、単語の分散表現を入力すると仮定します。 以下は入力を作るコードです。ハイパーパラメータについては適当に設定しています。

>>> from keras.layers import Input, Embedding, SimpleRNN, LSTM, GRU
>>> input = Input(shape=(100,))
>>> embedding = Embedding(input_dim=1000, output_dim=300)(input)

では、まずは SimpleRNN の内部状態を取得してみます。 return_state に True を設定することで、出力を一つ、内部状態を一つ返してきます。 ここで、出力は必ずリストの一番目に入っていることに注意してください。これは以降も同様です。

>>> SimpleRNN(units=100, return_state=True)(embedding)
[<tf.Tensor 'simple_rnn_4/TensorArrayReadV3:0' shape=(?, 100) dtype=float32>,
 <tf.Tensor 'simple_rnn_4/while/Exit_2:0' shape=(?, 100) dtype=float32>]

次に、LSTMの内部状態を取得してみます。 return_state に True を設定することで、出力一つ、内部状態2つを返してくることがわかります。

>>> LSTM(units=100, return_state=True)(embedding)
[<tf.Tensor 'lstm_6/TensorArrayReadV3:0' shape=(?, 100) dtype=float32>,
 <tf.Tensor 'lstm_6/while/Exit_2:0' shape=(?, 100) dtype=float32>,
 <tf.Tensor 'lstm_6/while/Exit_3:0' shape=(?, 100) dtype=float32>]

最後に、GRUの内部状態を取得してみます。 同じように return_state に True を設定することで、出力一つ、内部状態一つを返してきます。

>>> GRU(units=100, return_state=True)(embedding)
[<tf.Tensor 'gru_2/TensorArrayReadV3:0' shape=(?, 100) dtype=float32>,
 <tf.Tensor 'gru_2/while/Exit_2:0' shape=(?, 100) dtype=float32>]

課題と解決策

単純に RNN 単体で使う分には問題ないのですが、return_state=True を Bidirectional と合わせて使うと問題が起きます。 これは return_state を True にするとリストが返されるのですが、Bidirectional はリストを受け取ることを想定していないためです。 以下のようなエラーが発生します。

>>> Bidirectional(LSTM(units=100, return_state=True))(embedding)
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File ".../keras/engine/topology.py", line 602, in __call__
    output = self.call(inputs, **kwargs)
  File ".../keras/layers/wrappers.py", line 299, in call
    output = K.concatenate([y, y_rev])
  File ".../keras/backend/tensorflow_backend.py", line 1700, in concatenate
    rank = ndim(tensors[0])
  File ".../keras/backend/tensorflow_backend.py", line 544, in ndim
    dims = x.get_shape()._dims
AttributeError: 'list' object has no attribute 'get_shape'

対策として、手動で LSTM を2つ作成し、その結果を結合してあげる手法が考えられます。 その際、片方の LSTM は引数に go_backwards=True を渡します。 これにより、LSTM への入力が逆順に行われるようになります。

from keras.layers import LSTM, Bidirectional
from keras.layers.merge import Concatenate

fwd_state = LSTM(units=100, return_state=True)(embedding)[-1]
bwd_state = LSTM(units=100, return_state=True, go_backwards=True)(embedding)[-1]
state = Concatenate(axis=-1)([fwd_state, bwd_state])

参考資料

※ ちなみに、日本語のドキュメントには return_state 引数についての記載はありませんでした(2017/09/21時点)。