Tensorflowで使用するGPUおよびGPU台数とGPUメモリ使用率を指定する方法

何も考えずにGPUを積んだマシンでtensorflowを動かすとすべてのGPUを使ってしまう。これはこれで便利なのだが、他の人と共用のマシンだとリソース的に厳しい。そこで以下のようにして、Sessionに入る前に設定をする。

gpuConfig = tf.ConfigProto(
    gpu_options=tf.GPUOptions(per_process_gpu_memory_fraction=0.5),
    device_count={'GPU': 1})
with tf.Session(config=gpuConfig) as session:
    # 以下略

tf.GPUOptionsによりその設定をするわけだが、per_process_gpu_memory_fractionでGPUメモリの使用率を（×１００%）、device_countでGPUをいくつ使うかを指定する。上記の設定ではGPU１台を50%のメモリ使用率で用いるということである。どのGPUを使うかの設定は、自分はシェルの環境変数CUDA_VISIBLE_DEVICESによった。例えば以下のようにする。

ここでどの数字を与えるかでGPUの特定がなされるわけだが、どの数字を与えるべきかは環境依存である。何も設定せずに動かして全GPUをフルに使った場合、tensorflowのログにおそらく現れる使用GPUの通し番号が0から始まるという認識でよい。したがって、'0'を与えれば1台目のGPU、'1'を与えれば2台目のGPU、である。複数台ならexport CUDA_VISIBLE_DEVICES="0, 1"などとすればよい。