GPUが複数刺さった計算機でKeras(Tensorflow backend)を使おうとしたところ、
Tensorflowが全GPUのVRAMを全て確保しようとするために、他の計算が回らなくなって困った。
GPU1枚だけを専有するように指定する方法でハマったのでメモしておく。
ソリューション
環境変数のCUDA_VISIBLE_DEVICESで指定できる。
kerasやtensorflowをimportする前に宣言すること
import os
os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
os.environ["CUDA_VISIBLE_DEVICES"] = ""
例:nvidia-smiから見てGPU ID = 0のGPUを使用する場合は
os.environ["CUDA_VISIBLE_DEVICES"] = "0"
その他試したことなど
with文で with tf.device('/gpu:0'): などと書いた後で処理を記述する方法はうまくいかなかった。
KerasやTensorflowをimportした時点で全てのメモリが確保されてしまう。
環境
Ubuntu 16.04 LTS
Nvidia Driver Version: 384.98
CUDA 8.0
Keras 2.0.9
Tensorflow 1.4.0
GTX1080Ti x2
参考になった情報
[Can keras model run on specific device?](Can keras model run on specific device?)