メモリの少ないグラフィックボードで MNIST を動かす

先日、GeForce GT 730 を乗せた自作 PC の TensorFlow 周りの環境を整えました。

walkingmask.hatenablog.com

しかし、予想通り、環境構築に満足してしまって、肝心の TensorFlow をゴリゴリ書いたりはしてませんでした。

ところが、思いがけずインターンに参加させてもらえることになり、そこで TensorFlow を書く機会があったので、この環境を使ってみました。

その際に、低メモリグラボで TensorFlow を動かしていった方法を記録します。

ResourceExhaustedError

TensorFlow のチュートリアルの Deep MNIST for Experts を実行すると、ResourceExhaustedError が発生します。

f:id:walkingmask:20170429112057p:plain

これは、

accuracy.eval(feed_dict={x: mnist.test.images, y_: mnist.test.labels, keep_prob: 1.0})

この行が原因で起こります。

feed_dict の x に与えている mnist.test.images が、float32 型で 10000*28*28*32 個の要素を持つので、ざっと見積もっても 250880000*4 byte で、約 1 GB あるためです。

これでは、GT 730 のメモリサイズ 1 GB に、まず乗せられません。

また、グラボのメモリ上に先に乗せたモデルがすでに 812 MB 程度のメモリを占領しています。

sess.run(init)

をして、nvidia-smi を実行すると、実際にはトータルメモリサイズが 979 MiB (= 1026 MB) で、使用中のメモリが 802 MiB (= 840 MB) となっており、空きメモリが十分不足していることがわかります。

f:id:walkingmask:20170429112110p:plain

accuracy の feed_dict を分割

これを解消するためには、accuracy に与える feed_dict を分割して、得られた値をうまく整えてやれば良いと考えました。

まず、分割しない accuracy の計算は

correct_prediction = tf.equal(tf.argmax(y_conv, 1), tf.argmax(y_, 1))
accuracy = tf.reduce_mean(tf.cast(correct_prediction, tf.float32))

です。ここで行われていることは、1 行目で

モデルの出力と正解ラベルが同じであれば True に
入力データ (feed_dict: x) に対して行われるので戻り値は boolean の配列

2 行目で

これを tf.float32 にキャスト (True なら 1、False なら 0 に)
tf.reduce_mean で配列の平均値を取っている

なので、1 行目の操作は、個々のデータに対して行っても何ら変化はなく、2 行目の reduce_mean を分割して行えば良いわけです。

やってるのは、ただ平均を取ることなので、以下のよう for で総和してテストデータの数で割りました。

num_test = len(mnist.test.labels)
sum_accuracy = 0
for i in range(0, num_test, 50):
  sum_accuracy = sum_accuracy + accuracy.eval(feed_dict={x: mnist.test.images[i:i+50], y_: mnist.test.labels[i:i+50], keep_prob: 1.0})
print("test accuracy:", sum/(num_test/50))