tfdbg を使って Keras の nan や inf を潰す

はじめに

Keras に限らず、機械学習等の科学計算をおこなっているときに nan や inf が出現することがあります。
nan や inf は主にゼロ除算などのバグに起因して発生しますが、nan や inf それ自体を目的に使うこともあるため、エラーになるわけではありません。エラーになったとしても、「きっとどこかでゼロ除算しているんだろなぁ・・・」と思いを馳せることができるだけです。

そこで、TensorFlow 公式デバッガである tfdbg を使って、(TensorFlow バックエンドの) Keras モデルをデバッグして nan や inf の原因を追ってみます。

今回デモ用に利用したコードは公開していますので、ご興味があればご参照ください：
https://github.com/keisuke-nakata/keras-tfdbg-tutorial

tfdbg とは

公式の tfdbg ドキュメントのはじめには次のようにあります：

TensorFlow debugger (tfdbg) is a specialized debugger for TensorFlow. It lets you view the internal structure and states of running TensorFlow graphs during training and inference, which is difficult to debug with general-purpose debuggers such as Python's pdb due to TensorFlow's computation-graph paradigm.

抄訳すると、「tfdbg は TensorFlow 専用のデバッガで、Python の pdb みたいに TensorFlow をデバッグできるよ」という感じです。

TensorFlow をバックエンドとして Keras を利用している場合でも tfdbg することが出来ます。

今回対象とするモデル

今回デバッグの対象とするのは、 MNIST を2層のパーセプトロン (入力と出力を入れれば4層) で分類するようなモデルです：

_________________________________________________________________
Layer (type)                 Output Shape              Param #
=================================================================
dense_1 (Dense)              (None, 500)               392500
_________________________________________________________________
dense_2 (Dense)              (None, 10)                5010
=================================================================
Total params: 397,510
Trainable params: 397,510
Non-trainable params: 0
_________________________________________________________________

モデルは次のように構築しています (なお、以下では import やデータの読み込み等の本質ではない部分は省略していますので、詳細や再現に興味があれば元コードをご覧ください)：

def build_model():
    def reproducable_trunc_norm():
        return TruncatedNormal(mean=0.0, stddev=0.1, seed=RAND_SEED)

    def reproducable_const():
        return Constant(value=0.1)

    def unstable_categorical_crossentropy(y_true, y_pred):
        return - tf.reduce_sum(y_true * tf.log(y_pred), axis=len(y_pred.get_shape()) - 1)

    model = Sequential()
    model.add(Dense(
        HIDDEN_SIZE, input_shape=(FLAT_IMAGE_SIZE, ), activation='relu',
        kernel_initializer=reproducable_trunc_norm(), bias_initializer=reproducable_const()))
    model.add(Dense(
        NUM_LABELS, activation='softmax',
        kernel_initializer=reproducable_trunc_norm(), bias_initializer=reproducable_const()))

    model.compile(loss=unstable_categorical_crossentropy, optimizer=Adam(lr=0.025), metrics=['accuracy'])
    # model.compile(loss='categorical_crossentropy', optimizer=Adam(lr=0.025), metrics=['accuracy'])
    model.summary()
    return model

def reproducable_trunc_norm と def reproducable_const は、今回デバッグするにあたって挙動を固定するための乱数調整です。
また、今回 nan を引き起こす原因となっているのは def unstable_categorical_crossentropy を model の loss に指定しているためです。今回はデモなので事前に原因が明らかですが、気づかなかったものとして話を進めます。正しくは loss='categorical_crossentropy' ですね (コメントアウトされている行)。

では、モデルを学習させてみましょう。

x_train, y_train, x_test, y_test = get_mnist_dataset()
model = build_model()
model.fit(x_train, y_train, batch_size=100, epochs=10, validation_data=(x_test, y_test))

loss が nan ですね。つらいです。
(上の結果は、python debug_mnist.py とすると得られます)

デバッグしてみる

tfdbg を起動してみる

早速デバッグしてみましょう。
まずは Keras が裏で使用している TensorFlow のセッションをハックしにいって、デバッガを仕込みます。以下の関数を最初に呼べば OK です¹。

def set_debugger_session():
    sess = K.get_session()
    sess = tf_debug.LocalCLIDebugWrapperSession(sess)
    sess.add_tensor_filter('has_inf_or_nan', has_inf_or_nan)
    K.set_session(sess)

python debug_mnist.py --debug を実行すると、set_debugger_session() を最初に呼んでから Keras のモデルの学習が始まるようにプログラムを書いているため、次のような画面が立ち上がると思います：

無事 tfdbg が起動しました！
おもむろに run コマンドを二回入力²してから、 lt (list_tensors の省略形です) を打ってみてください：

TensorFlow のグラフ上に存在するテンソルの一覧が表示されました！
もちろんテンソルの中身を確認することも出来ます。
pt dense_1/kernel:0 を入力してみてください。(pt は print_tensor の省略形です)
(なお、lt で表示したテンソルの名前をおもむろにクリックしても、そのテンソルを pt したのと同じことをしてくれます)

いま表示しているのは dense_1/kernel:0 ですが、これは、先程の model の構造と照らし合わせると、「第1層目の Dense レイヤーの kernel」、すなわち重み行列 $W$ が表示されています。(Keras では学習対象となる主パラメータのことを kernel と呼ぶことが多いです)
shape: (784, 500) というのも、MNIST の画像サイズが 28*28=784 で、最初の隠れ層のノー数を 500 に設定していることと辻褄が合っています。

nan と inf を探す

tfdbg はデバッガなので、「この条件が満たされるまで実行」もできます。
(そのために session に登録した関数が has_inf_or_nan です。これは自分で作成して自由に名前をつけて登録することが出来ます)。

nan や inf の原因を探すには、実際に nan や inf が出現するまで何回か run を実行しないといけません。上述の機能を使って、モデルのどこかに nan または inf が出現するまで run を実行してみましょう。

run -f has_inf_or_nan

tfdbg のタイトルバーに「run-end: run #2: 3 fetches; 3 feeds」と表示され、実行が停止しました。どうやら nan か inf が出現したようです。
lt を実行し、テンソル一覧が表示されることを確認してください。

さらに、lt -f has_inf_or_nan を実行することで、「どのテンソルが nan または inf を持っているのか」を表示することが可能です。

問題のテンソルがひとつしかありません。中身を表示してみます pt Const_4:0

こいつの中身は array(inf, dtype=float32) でした。確かに inf です。
もう1回だけ実行を進めてみます。run を実行してから、lt -f has_inf_or_nan を実行してください：

今度は Const_4:0 だけではなく、たくさんテンソルの中身に nan または inf が入ってしまっているようです。

さて、テンソル一覧の上のほうを見ると、Log:0, mul:0, Sum:0, Neg:0 等が並んでいますね。
この計算グラフの並びは、もう明らかに自分で定義したロス関数である

    def unstable_categorical_crossentropy(y_true, y_pred):
        return - tf.reduce_sum(y_true * tf.log(y_pred), axis=len(y_pred.get_shape()) - 1)

のことです (最初に tf.log(y_pred) (つまり Log:0) が計算され、次に * (つまり mul:0) が、次に reduce_sum (つまり Sum:0) が、次に - (つまり Neg:0) が計算されていますね)。
では Log:0 の中身を見てみましょう。pt Log:0：

-inf がいます！見づらければ、おもむろに /inf とでも打つとハイライトしてくます。もう Log の時点で -inf が来ている時点でアウトですね。ログの結果が -inf になるのは、$\log(0)$ が評価されたからに違いありません。
上のメニューから、list_inputs をクリックしてみてください：

Log に対する入力が表示されます。最後の活性化関数である Softmax が見えますね。pt dense_2/Softmax:0 を実行してみてください：

思った通り、Log で -inf になっていた部分の値が 0 です。
これでは loss の値が nan になるのも当然です。

ここまでくれば、あとは「softmax crossentropy inf」とかでググれば問題はほぼ解決したようなものです。
今回の例なら、きっと「softmax と crossentropy を組み合わせて使うときは数値計算の都合上、特殊なことをしないといけないから、自分でロス定義しないでライブラリの実装を使いなさい」的な情報が見つかることでしょう。

tfdbg を終了するときは、 quit コマンドを実行してください。

原因がわかったので、コードのコメントアウトされている部分を復活させて実行すると、無事 nan を回避することができるようになります：

$ python debug_mnist.py
Using TensorFlow backend.
_________________________________________________________________
Layer (type)                 Output Shape              Param #
=================================================================
dense_1 (Dense)              (None, 500)               392500
_________________________________________________________________
dense_2 (Dense)              (None, 10)                5010
=================================================================
Total params: 397,510
Trainable params: 397,510
Non-trainable params: 0
_________________________________________________________________
Train on 60000 samples, validate on 10000 samples
Epoch 1/10
2017-11-26 01:35:25.856563: W tensorflow/core/platform/cpu_feature_guard.cc:45] The TensorFlow library wasn't compiled to use SSE4.1 instructions, but these are available on your machine and could speed up CPU computations.
2017-11-26 01:35:25.856604: W tensorflow/core/platform/cpu_feature_guard.cc:45] The TensorFlow library wasn't compiled to use SSE4.2 instructions, but these are available on your machine and could speed up CPU computations.
2017-11-26 01:35:25.856621: W tensorflow/core/platform/cpu_feature_guard.cc:45] The TensorFlow library wasn't compiled to use AVX instructions, but these are available on your machine and could speed up CPU computations.
2017-11-26 01:35:25.856629: W tensorflow/core/platform/cpu_feature_guard.cc:45] The TensorFlow library wasn't compiled to use AVX2 instructions, but these are available on your machine and could speed up CPU computations.
2017-11-26 01:35:25.856638: W tensorflow/core/platform/cpu_feature_guard.cc:45] The TensorFlow library wasn't compiled to use FMA instructions, but these are available on your machine and could speed up CPU computations.
60000/60000 [==============================] - 14s - loss: 1.6388 - acc: 0.8431 - val_loss: 0.1992 - val_acc: 0.9450
Epoch 2/10
60000/60000 [==============================] - 8s - loss: 0.1921 - acc: 0.9482 - val_loss: 0.2155 - val_acc: 0.9468
Epoch 3/10
60000/60000 [==============================] - 10s - loss: 0.1814 - acc: 0.9529 - val_loss: 0.2108 - val_acc: 0.9477
Epoch 4/10
60000/60000 [==============================] - 12s - loss: 0.1734 - acc: 0.9567 - val_loss: 0.2226 - val_acc: 0.9507
Epoch 5/10
60000/60000 [==============================] - 14s - loss: 0.1728 - acc: 0.9584 - val_loss: 0.2229 - val_acc: 0.9516
Epoch 6/10
60000/60000 [==============================] - 24s - loss: 0.1670 - acc: 0.9608 - val_loss: 0.2424 - val_acc: 0.9566
Epoch 7/10
60000/60000 [==============================] - 17s - loss: 0.1707 - acc: 0.9609 - val_loss: 0.2184 - val_acc: 0.9565
Epoch 8/10
60000/60000 [==============================] - 13s - loss: 0.1697 - acc: 0.9626 - val_loss: 0.2641 - val_acc: 0.9514
Epoch 9/10
60000/60000 [==============================] - 7s - loss: 0.1619 - acc: 0.9637 - val_loss: 0.2136 - val_acc: 0.9631
Epoch 10/10
60000/60000 [==============================] - 13s - loss: 0.1603 - acc: 0.9644 - val_loss: 0.2962 - val_acc: 0.9505
Test loss: 0.296182284909
Test accuracy: 0.9505

HAPPY

ちなみに

softmax の値が 0 になるのは、softmax に入力されるベクトルの中身に巨大な数値が含まれていて、softmax の式の分母の基準化が効きすぎるときに起こります。exp を取るので、数値の大きさが超拡大されて情報落ち³が発生するからです。ただし、softmax の出力が 0 になる事自体はおかしなことではありません (そのクラスに所属する確率が 0 であるというのは有効な予測です)。問題なのは、その後に log を取っているので log(0) が評価されてしまうことです。

softmax は入力に定数を足し引きしても結果が変わらないという性質があるので、それを使って入力の最大値でスケーリングをかけることで情報落ちを回避できます。
また、cross_entropy と softmax は超簡単に言えばそれぞれ log と exp を取る演算なので、$\text{crossentropy}(\text{softmax}(x))$ を計算するのではなく、まず softmax と crossentropy の合成関数を作って等価に変形してやると一部の log と exp を打ち消すことができます。
以上の工夫をすると、 softmax の結果がほぼゼロになるときでも -inf にならずにすみます。

よって、TensorFlow や chainer では、「softmax をとってから cross_entropy を計算」するのではなく、「softmax と cross_entropy を同時に計算」する関数が用意されています。
TensorFlow なら tf.nn.softmax_cross_entropy_with_logits ですし、chainer なら chainer.functions.softmax_cross_entropy です。

一方で、Keras では、「softmax をとってから cross_entropy を計算」します。話が違う、と思うかもしれませんが、Keras は log(0) を防ぐために softmax からの出力を $[\epsilon, 1 - \epsilon]$ の領域に寄せてから cross_entropy を計算して inf 問題を回避しています。 (https://github.com/fchollet/keras/blob/master/keras/backend/tensorflow_backend.py#L2888)

参考資料

• デバッガを使ってKerasのモデルをデバッグする - 自然言語処理の深遠
• Debugging TensorFlow Programs | TensorFlow
  • TensorFlow 公式ドキュメント
• Google Developers Japan: tfdbg による TensorFlow モデルのデバッグ
  • tfdbg を紹介する Google Developers Blog (Japan) の記事
• tensorflow/debug_mnist.py at master · tensorflow/tensorflow
  • TensorFlow 公式の tfdbg のサンプル
• keras/mnist_cnn.py at master · fchollet/keras
  • Keras の MNIST サンプル
• keras/tensorflow_backend.py at master · fchollet/keras
  • Keras が catgorical_crossentropy の不安定性をどうやって解消しているかの部分

1. has_inf_or_nan は、tfdbg を TensorFlow から直接呼んだときは最初からセットされているようですが、Keras から呼ぶとなぜかセットされていなかったので、自分で設定しています。 ↩

2. 二回入力しているのは、テンソルを初期化させるためです。run コマンド一回が TensorFlow 側の session.run() 1回に対応するのですが、Keras を使っていると隠蔽しまいます。おそらく1回目の run ではテンソルの定義のみがおこなわれているようです。 ↩

3. exp(1e-3)/exp(1e+3) とか計算したら分母がでかすぎて結果がゼロになってしまう ↩