終末 A.I.

データいじりや機械学習するエンジニアのブログ

TensorFlow の ptb サンプルを動かす

TensorFlow で ptb を学習させるサンプルは Recurrent Neural Networks チュートリアルに記載されている通り、 github からソースを落としてきて、tensorflow/models/rnn/ptb に移動し、ptb_word_lm.py を動かすだけで簡単に動作させることができます。

ただし残念なことに、サンプルコードをぱっと見ただけで使い方を理解するのは相変わらずなかなか難しいです。API にのっていないクラスや関数が使われたり、RNN 特有の処理がさくっと書かれているため、サンプルとしてまあ不親切なわけです。

そこで今回も前回の DNN と同様、自己流にサンプルを書き直して使い方を確認してみました。書き直したコードを動作させてみて、だいたい元のコードと同じような性能が出ていそうなので、ざっくりと解説をしてみようと思います。

いつもどおり完成コードは、記事の最後の方に置かせていただいています。書き直したコード内の内部で reader ファイルは元サンプルの reader と全く同じものを利用しています。

また、RNNをまともに学習させようと思うとCPUでは非常に遅いので、動作確認には AWSGPU インスタンスを使用しています。

TensorFlow ってどうやって使うの?AWSGPUインスタンスでどうやって動かせばいいの?という方は、まずは過去記事をご参照ください。

ハイパーパラメーターと入出力用変数

さて、手法のコードを書く前にイニシャライズをする必要がありますが、DNN の場合だいたいここでハイパーパラメーターの設定と入出力用のデータを定義します。TensorFlow で書く場合もその流儀にならって書くと下記のようになります。

# config
self._batch_size = 20
self._num_steps = 2
self._hidden_size = 2
self._vocab_size = 10000
self._num_layers = 1
self._keep_prob = 0.5
self._max_grad_norm = 1

# input and output variables
self._input_data = tf.placeholder(tf.int32, [self._batch_size, self._num_steps])
self._targets = tf.placeholder(tf.int32, [self._batch_size, self._num_steps])
self._initial_state = None
self._final_state = None
self._cost = None
self._train_op = None
self._logits = None

バッチサイズなどがハイパーパラメーターで、input_data などが PlaceholderVariable などの入出力用の変数となります。Placeholder も Variable も TensorFlow で計算の時に使用される変数クラスですが、Placeholder が入力用、Variable は最終出力および中間出力用に主に使用されるものでしたね。

RNN で大事なものは、内部状態をいつまで保持するかを定義している num_steps と内部状態を表す変数である initial_state と final_state です。この2つの値に注目して以降のコードを見ていきましょう。

LSTMの利用

続いて、LSTMを利用して、次の単語を出力させる手順までを見てみましょう。損失関数の定義やパラメーター更新の方法は通常のDNNと大差ありませんので省略させていただきます。

# LSTM
lstm_cell = tf.nn.rnn_cell.BasicLSTMCell(self._hidden_size, forget_bias=0.0, state_is_tuple=True)
# add dropout
if is_training:
    lstm_cell = tf.nn.rnn_cell.DropoutWrapper(lstm_cell, output_keep_prob=self._keep_prob)
# add multi lyaers
cell = tf.nn.rnn_cell.MultiRNNCell([lstm_cell] * self._num_layers, state_is_tuple=True)
# initial state setup
self._initial_state = cell.zero_state(self._batch_size, tf.float32)

ここから先のコードは色々と不親切で全然サンプルになっておりませんので心して追っかけていきましょう。

まず、LSTM の基本クラスである tf.nn.rnn_cell.BasicLSTMCell が早速API一覧にのっていません。詳細が気になる方は、TensorFlow ソースコード内の tensorflow/python/ops/rnn_cell.py をご参照ください。使い方自体はそんなに難しくはなく、隠れ変数のサイズと忘却率、そして state_is_tuple という値を指定しているだけで LSTM セルを定義できてしまいます。

state_is_tuple は元のサンプルコードでは使用されていませんが、使用しないままコードを動かすと、state_is_tuple が False だと遅いし False の場合はそのうち Duplicated だからね!という警告 が r0.9 の時点では表示されます。このオプションは、LSTM の計算時に入出力される隠れ変数を Variable クラスのタプルとして表現するかどうかを示すものになります。

書き換えるだけで早くなるのであれば書き換えない手はありませんが、evaluation の時に独特な処理が必要となりますのでそこだけご注意ください。下記のコードの run_epoch 関数を参考にするとわかりやすく汎用的に書き換えることができますのでご参照ください。

続いて用いられているのが、tf.nn.rnn_cell.DropoutWrapper と tf.nn.rnn_cell.MultiRNNCell です。この2つは見たまんまなのですが、LSTM の入力にドロップアウトを適用する処理と、LSTM + Dropout を複数個並べてくれる処理になります。

使用している引数も、DropoutWrapper の場合はベースとなるLSTMセルとキープレート、MultiRNNCellの場合は組み合わせるセルの配列と上記に登場した state_is_tuple を指定しています。

最後に、RNNCell クラスの zero_state という便利関数を利用して、隠れ変数の初期値を表す Placeholder クラスを取得します。この値はのちのちの計算処理で必要になるので確実に保持しておく必要があります。

# Load predefined layer "embedding"
with tf.device("/cpu:0"):
    embedding = tf.get_variable("embedding", [self._vocab_size, self._hidden_size])
    inputs = tf.nn.embedding_lookup(embedding, self._input_data)

# Add dropout after embedding layer
if is_training:
    inputs = tf.nn.dropout(inputs, self._keep_prob)

次に、単語IDを単語ベクトルに変換する処理を行います。ここでもなかなか TensorFlow を作っている人でないと普通わからないだろうという処理を行っていまして、tf.get_variable で "embedding" という変数を取得しています。どうやら "embedding" という変数が Word Embedding 用にすでに用意されているらしく、その変数を使って単語IDをベクトルにする処理を行っているようです。仕様書読むだけでは分かりませんね。

ちなみに rnn_cell.py には入力を単語ベクトルにした後、RNN セルに適用してくれる EmbeddingWrapper クラスが存在しています。試してみていませんが多分同じような結果になると思われます。

# Calculate LSTM Layer for in _num_steps
outputs = []
state = self._initial_state
with tf.variable_scope("RNN"):
    for time_step in range(self._num_steps):
        if time_step > 0: tf.get_variable_scope().reuse_variables()
        (cell_output, state) = cell(inputs[:, time_step, :], state)
        outputs.append(cell_output)

次に、LSTM を利用した計算の肝となるシーケンスの処理部分を見てみましょう。今回のインプットデータは、ミニバッチ × シーケンス長(num_steps) × 単語ベクトル という形の配列として作成されています。ミニバッチ単位でベクトルを計算していく点は他の DNN と変わりありませんが、RNN ではシーケンシャルなデータの処理を行う必要があります。

LSTM のような隠れ変数がある層では、truncate という隠れ変数をリセットし、バックプロバケーションを打ち切る長さを決めておく必要があります。この長さはもちろん可変でも良いのですが、それではミニバッチの利用に向かず処理が遅くなるため、このサンプルでは固定長とされています。

さてこの truncate 処理ですが、それ用の関数などはなく、上記のように自前で実装する必要があります。とはいっても見ての通りで num_steps 数分、一個ずつ入力データをずらして RNN セルに入力していっているだけです。注意点としては、後の評価やパラメーター更新に使用するため、出力される変数はなんらかの形で保持しおく必要がある点くらいです。

# Final output layer for getting word label 
output = tf.reshape(tf.concat(1, outputs), [-1, self._hidden_size])
softmax_w = tf.get_variable("softmax_w", [self._hidden_size, self._vocab_size])
softmax_b = tf.get_variable("softmax_b", [self._vocab_size])
self._logits = tf.matmul(output, softmax_w) + softmax_b

最後に、出力を単語ID に変換するための通常の全結合層を挿入して終了です。先ほど取得した num_steps 数分の出力を全て処理してやらなければならない点がややこしいですが、注意すべきはそこくらいとなります。このあと損失関数を定義し、最適化関数をはさんで Session.run を実行すれば学習は完了です。

予測された次の単語を取得する処理も、logits をミニバッチおよびステップ毎に argmax などで候補となる単語IDに変換してやればOKです。rnn_cell.py の中身を熟読すればもっとシンプルに書けそうですね。

TensorFlow basic RNN sample