CUDA_ERROR_OUT_OF_MEMORYになる

tensorflow用のgpuマシンで学習をさせようと、早速大量の画像を食わせたら、長い間画像を読んだ後、

CUDA_ERROR_OUT_OF_MEMORY; total memory reported:

とエラーが出た。

tensorflowのGPU版では、デフォルトではマシンにのっている全GPUの全メモリを使用する。
そこで使用するGPUを制限させることにした。

次のコードを追加。

    config = tf.ConfigProto(
      gpu_options=tf.GPUOptions(
          per_process_gpu_memory_fraction=0.8, # 最大値の80%まで
          allow_growth=True # True->必要になったら確保, False->全部
      )
    )
    sess = sess = tf.Session(config=config)

そしたらちゃんと動いた。

ちなみに、GPUを使わないようにするには、次のコードを入れればよい。
自分はGPUで回しているのでコメントアウトされている。

#    config = tf.ConfigProto(
#      device_count = {'GPU': 0}
#    )
#    sess = tf.Session(config=config)