Kerasメモリ制限

KerasでbackendにTensorlfowを使っていると、GPUメモリを全て食ってしまう。 Kerasのバックエンドに使っているTensoflowに設定を加える。

tf.ConfigProtoで設定を加えたら好きなように制限できる。

以下の設定だと必要な分だけ確保してくれる。

import tensorlfow as tf
from keras.backend import tensorflow_backend

gpuConfig = tf.ConfigProto(allow_soft_placement=True,
                           gpu_options=tf.GPUOptions(allow_growth=True,
　　　　　　　　　　　　　　　　　　　　                   visible_device_list="0"))

sess = tf.Session(config=gpuConfig)
tensorflow_backend.set_session(sess)

GPUの割合で指定したかったら以下のような感じ。

gpuConfig = tf.ConfigProto()
gpuConfig.gpu_options.visible_device_list = "0, 1"
gpuConfig.gpu_options.per_process_gpu_memory_fraction = 0.4

tf.kerasを使うときは

import tensorlfow as tf
from tensorflow.keras import backend

gpuConfig = tf.ConfigProto(allow_soft_placement=True,
                           gpu_options=tf.GPUOptions(allow_growth=True,
　　　　　　　　　　　　　　　　　　　　                   visible_device_list="0"))

sess = tf.Session(config=gpuConfig)
backend.set_session(sess)