Python: Keras/TensorFlow で GPU のメモリを必要な分だけ確保する

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のより良い書き方を学ぶ