Keras のバックエンドに TensorFlow を使う場合、デフォルトでは一つのプロセスが GPU のメモリを全て使ってしまう。 今回は、その挙動を変更して使う分だけ確保させるように改めるやり方を書く。
環境には次のようにしてセットアップした Ubuntu 16.04 LTS を使っている。 blog.amedama.jp
サンプルとして動作させるアプリケーションには Keras が提供している MNIST データセットを CNN で認識するものを使う。 まずはこれをダウンロードしておこう。 同時に、セッションをクリアするパッチも追加しておく。
$ wget https://raw.githubusercontent.com/fchollet/keras/master/examples/mnist_cnn.py $ echo 'K.clear_session()' >> mnist_cnn.py
上記を実行すると GPU を使ったニューラルネットワークの学習が始まる。
$ python mnist_cnn.py
学習している最中に、別のターミナルから nvidia-smi
コマンドを実行してみよう。
すると、ビデオカードに載っているメモリのほとんど全てを上記のプロセスが使っていることが分かる。
$ nvidia-smi Wed Jun 7 21:28:52 2017 +-----------------------------------------------------------------------------+ | NVIDIA-SMI 375.66 Driver Version: 375.66 | |-------------------------------+----------------------+----------------------+ | GPU Name Persistence-M| Bus-Id Disp.A | Volatile Uncorr. ECC | | Fan Temp Perf Pwr:Usage/Cap| Memory-Usage | GPU-Util Compute M. | |===============================+======================+======================| | 0 GeForce GTX 105... Off | 0000:01:00.0 Off | N/A | | 49% 64C P0 63W / 75W | 3863MiB / 4038MiB | 87% Default | +-------------------------------+----------------------+----------------------+ +-----------------------------------------------------------------------------+ | Processes: GPU Memory | | GPU PID Type Process name Usage | |=============================================================================| | 0 1874 C python 3861MiB | +-----------------------------------------------------------------------------+
全部で 4GB しかないメモリのうち 3.8GB を一つのプロセスが使っている。
この状況で、別のターミナルからもう一つ Keras のプロセスを動かしてみよう。 当たり前だけど、これは残りのメモリが少なすぎて実行に失敗する。
$ python mnist_cnn.py ...(snip)... 2017-06-07 21:46:15.514867: E tensorflow/stream_executor/cuda/cuda_driver.cc:893] failed to allocate 134.44M (140967936 bytes) from device: CUDA_ERROR_OUT_OF_MEMORY 2017-06-07 21:46:15.835973: E tensorflow/stream_executor/cuda/cuda_dnn.cc:359] could not create cudnn handle: CUDNN_STATUS_INTERNAL_ERROR 2017-06-07 21:46:15.836015: E tensorflow/stream_executor/cuda/cuda_dnn.cc:326] could not destroy cudnn handle: CUDNN_STATUS_BAD_PARAM 2017-06-07 21:46:15.836032: F tensorflow/core/kernels/conv_ops.cc:659] Check failed: stream->parent()->GetConvolveAlgorithms(&algorithms) Aborted (core dumped)
これでは一つのマシンで同時に学習させられるモデルが一つだけになってしまう。
この挙動を変更するには Keras が使う TensorFlow のセッションの設定を変更する必要がある。
TensorFlow には GPU のオプションとして allow_growth
というものがあり、これを有効にすると必要な分だけ確保するようになる。
あとは、そう設定した TensorFlow のセッションを Keras で使うようにできれば上手くいく。
これには keras.backend.tensorflow_backend
モジュールにある set_session()
という関数を使う。
その部分だけをスニペットにすると、こんな感じ。
import tensorflow as tf from keras.backend import tensorflow_backend config = tf.ConfigProto(gpu_options=tf.GPUOptions(allow_growth=True)) session = tf.Session(config=config) tensorflow_backend.set_session(session)
試しに、さっきの MNIST サンプルのコードに上記を組み込んで動作を確認してみよう。
$ cat << 'EOF' > /tmp/keras-mnist.patch import tensorflow as tf from keras.backend import tensorflow_backend config = tf.ConfigProto(gpu_options=tf.GPUOptions(allow_growth=True)) session = tf.Session(config=config) tensorflow_backend.set_session(session) EOF $ sed -i -e ' /^from keras import backend as K$/r /tmp/keras-mnist.patch ' mnist_cnn.py
上記を実行すると、こんな感じで挿入される。 もちろん自分でエディタを使って編集しても構わない。
$ head -n 22 mnist_cnn.py | tail -n 10 from keras.layers import Conv2D, MaxPooling2D from keras import backend as K import tensorflow as tf from keras.backend import tensorflow_backend config = tf.ConfigProto(gpu_options=tf.GPUOptions(allow_growth=True)) session = tf.Session(config=config) tensorflow_backend.set_session(session) batch_size = 128
再度実行してみよう。
$ python mnist_cnn.py
そして別のターミナルから nvidia-smi
コマンドを叩いてメモリの消費量を確認する。
$ nvidia-smi Wed Jun 7 21:32:04 2017 +-----------------------------------------------------------------------------+ | NVIDIA-SMI 375.66 Driver Version: 375.66 | |-------------------------------+----------------------+----------------------+ | GPU Name Persistence-M| Bus-Id Disp.A | Volatile Uncorr. ECC | | Fan Temp Perf Pwr:Usage/Cap| Memory-Usage | GPU-Util Compute M. | |===============================+======================+======================| | 0 GeForce GTX 105... Off | 0000:01:00.0 Off | N/A | | 44% 59C P0 63W / 75W | 425MiB / 4038MiB | 87% Default | +-------------------------------+----------------------+----------------------+ +-----------------------------------------------------------------------------+ | Processes: GPU Memory | | GPU PID Type Process name Usage | |=============================================================================| | 0 1914 C python 423MiB | +-----------------------------------------------------------------------------+
さっきは 3.8GB も使っていたけど、今度は 423MB しか使っていない!
これなら複数のモデルを同時に学習させることができる。
$ nvidia-smi Wed Jun 7 21:58:02 2017 +-----------------------------------------------------------------------------+ | NVIDIA-SMI 375.66 Driver Version: 375.66 | |-------------------------------+----------------------+----------------------+ | GPU Name Persistence-M| Bus-Id Disp.A | Volatile Uncorr. ECC | | Fan Temp Perf Pwr:Usage/Cap| Memory-Usage | GPU-Util Compute M. | |===============================+======================+======================| | 0 GeForce GTX 105... Off | 0000:01:00.0 Off | N/A | | 41% 55C P0 66W / 75W | 848MiB / 4038MiB | 98% Default | +-------------------------------+----------------------+----------------------+ +-----------------------------------------------------------------------------+ | Processes: GPU Memory | | GPU PID Type Process name Usage | |=============================================================================| | 0 2315 C python 423MiB | | 0 2348 C python 423MiB | +-----------------------------------------------------------------------------+
めでたしめでたし。
ところで、逐次的にメモリを確保するとなるとパフォーマンスに影響がないかが気になる。 そこで、編集前と編集後で学習にかかる時間に変化があるかについても調べてみた。
まずは編集前から。
$ time python mnist_cnn.py
...(snip)...
real 2m4.539s
user 1m59.228s
sys 0m11.508s
2 分 4 秒で終わっている。
次は編集後を。
$ time python mnist_cnn.py
...(snip)...
real 2m4.666s
user 1m59.800s
sys 0m11.480s
こちらも 2 分 4 秒で終わった。 どうやらパフォーマンスに大きな影響は無さそうだ。
まとめ
今回は Keras のバックエンドを TensorFlow で動かすときに必要な分だけメモリを確保するやり方について書いた。
スマートPythonプログラミング: Pythonのより良い書き方を学ぶ
- 作者: もみじあめ
- 発売日: 2016/03/12
- メディア: Kindle版
- この商品を含むブログ (1件) を見る