tensorflow

tf.kerasでメモリを必要分だけ確保する設定

設定せずにGPUを使用すると、一つのプロセスでGPUのメモリをほぼマックス確保してしまいます。

そこまで必要でないことが多く、またその場合は1GPUで複数のプロセスを動かしたいので、メモリも必要な分だけ確保するように設定します。

tensorflowでは、allow_growthというオプションでこの設定ができます。
GPU使用に関する公式ドキュメント

本家のkerasの場合、設定方法はgithubのissueに書いてありました。

本家のkerasの場合(上のissueのコピペです)

import tensorflow as tf
from keras.backend.tensorflow_backend import set_session
config = tf.ConfigProto()
config.gpu_options.allow_growth = True  # dynamically grow the memory used on the GPU
config.log_device_placement = True  # to log device placement (on which device the operation ran)
                                    # (nothing gets printed in Jupyter, only if you run it standalone)
sess = tf.Session(config=config)
set_session(sess)  # set this TensorFlow session as the default session for Keras

tf.kerasを使用している場合は、set_sessionが提供されているモジュール名が違います。

tensorflow内のapiなので、tensorflow_backendなんていうモジュールは無いようで、tensorflow.keras.backend.set_sessとなります。
公式APIドキュメント

allow_growthを設定するコードは以下のようになります。

import tensorflow as tf
from tf.keras.backend import set_session
config = tf.ConfigProto()
config.gpu_options.allow_growth = True  # dynamically grow the memory used on the GPU
config.log_device_placement = True  # to log device placement (on which device the operation ran)
                                    # (nothing gets printed in Jupyter, only if you run it standalone)
sess = tf.Session(config=config)
set_session(sess)  # set this TensorFlow session as the default session for Keras

コメントを残す

メールアドレスが公開されることはありません。 * が付いている欄は必須項目です