Tensorflowのメモリ確保をChainerっぽくする #DeepLearning

Tensorflow (v0.12時点) はデフォルトではCUDAから見えている全GPUのメモリを専有してしまいます。また、GPUを使い学習したモデルをCPUで使うことができません。これを、Chainerと同じように、

with device: ...のような形で簡単にGPU/CPUを切り替え
メモリ使用にあわせて自動的にGPUメモリを増減

することを考えます。

そのような設定にするためには、tf.Sessionを下記のように作ります。

config = tf.ConfigProto(allow_soft_placement=True)
config.gpu_options.allow_growth = True

with tf.Session(config=config) as sess:
    with tf.device(my_device_option):
        model = MyGreatModel()
    sess.run(tf.global_variables_initializer())
    ...

tf.ConfigProtoには他にもオプションがありますが、それは他に詳しいです。