Edited at

マルチGPU環境で、計算に使用するCUDAデバイスを選択する

More than 1 year has passed since last update.

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?