KerasでGPUメモリの使用量を抑える方法

バックエンドをTensorFlowとしてKerasを利用しようとすると，デフォルトだとGPUのメモリを全部使う設定になっていて複数の実験を走らせられないので，GPUメモリの使用量を抑える設定方法について紹介します．^{1 2}

検証環境

tensorflow==1.3.0
tensorflow-gpu==1.3.0
Keras==2.0.6

GPUのメモリ使用量を抑える

下記コードを貼り付けるなり，importするなりすれば大丈夫です．

最小限のGPUメモリのみ確保

gpu_options.allow_growthで設定可能です．
実行時に必要な分だけ確保する方法で，さらに必要になるとメモリ領域を拡張するようになっています．
ただし自動的にはメモリを開放しないので，メモリが断片化して性能が悪化する可能性があり注意が必要です．¹

import tensorflow as tf
from keras import backend as K
config = tf.ConfigProto()
config.gpu_options.allow_growth = True
sess = tf.Session(config=config)
K.set_session(sess)

合計GPUメモリの割合指定で確保

gpu_options.per_process_gpu_memory_fractionで設定可能です．
下記例だと40%のメモリを使います．

import tensorflow as tf
from keras import backend as K
config = tf.ConfigProto()
config.gpu_options.per_process_gpu_memory_fraction = 0.4
sess = tf.Session(config=config)
K.set_session(sess)

実行時の状況

$ nvidia-smi -l
+-----------------------------------------------------------------------------+
| Processes:                                                       GPU Memory |
|  GPU       PID  Type  Process name                               Usage      |
|=============================================================================|
|    0     30654    C   python                                        3527MiB |
|    1     30779    C   python                                        3357MiB |
+-----------------------------------------------------------------------------+

References

1. TensorFlow - Using GPUs ↩ ↩²

2. Keras FAQ: Frequently Asked Keras Questions ↩