Tensorflow系のライブラリで使用するGPUを指定する

はじめに

Tensorflowでは，デフォルトの設定で回すと搭載されているGPUすべてのメモリを確保してしまいます．TensorflowをバックエンドにしたKerasも同様です．個人のデスクトップPCで回す場合には問題ありませんが，複数のGPUを搭載した共有のGPUサーバを利用する場合，大変迷惑です．
以下に使用するGPUを指定するサンプルコードをまとめました．

Tensorflow

import tensorflow as tf
config = tf.ConfigProto(
    gpu_options=tf.GPUOptions(
        visible_device_list="0,1", # specify GPU number
        allow_growth=True
    )
)
sess = tf.Session(config=config)

Sessionを確保するときにconfigを渡してやります．
allow_growth=Trueを指定することによって，指定したGPUのメモリを最初にすべて確保するのではなく，必要に応じて確保するように設定できます．
あまりうれしくはありませんが他に空いているGPUがない場合，メモリが空いていれば複数のプログラムを回すことが可能です．

keras

import tensorflow as tf
from keras.backend.tensorflow_backend import set_session
config = tf.ConfigProto(
    gpu_options=tf.GPUOptions(
        visible_device_list="2", # specify GPU number
        allow_growth=True
    )
)
set_session(tf.Session(config=config))

(https://github.com/fchollet/keras/issues/1538)

基本的にはバックエンドとして動作するTensorflowの設定を変更します．keras.backend.tensorflow_backend.set_sessionを用います．
注意点としては，コードの内部でtf.Sessionを取得しているようなケースでは，先頭に上記を記述しても意味がありません（とくに自分で書いていないコードを走らせる場合）．

(nvidia-docker)

nvidia-dockerにはコンテナに対してマウントするGPUをを設定するオプションがあります．

$ NV_GPU=3,4 nvidia-docker run HOGEHOGE

NV_GPUによって設定できます．