TensorflowやKerasでGPUメモリの使用量をコマンドラインから制限する

Tensorflowや、バックエンドにTensorflowを使うKerasは、プロセスが実行された時にGPUのメモリを上限一杯まで確保してしまう。

以下のプログラムをpythonファイルに書き込めばGPUを制限できるが、GPUメモリを全部使って実行したい時もあるし、したくない時もある。

gpus = tf.config.experimental.list_physical_devices('GPU')
if gpus:
  try:
    # Currently, memory growth needs to be the same across GPUs
    for gpu in gpus:
      tf.config.experimental.set_memory_growth(gpu, True)
    logical_gpus = tf.config.experimental.list_logical_devices('GPU')
    print(len(gpus), "Physical GPUs,", len(logical_gpus), "Logical GPUs")
  except RuntimeError as e:
    # Memory growth must be set before GPUs have been initialized
    print(e)

とはいえ毎回全部のpythonファイルに書くのも面倒くさいし、コメントアウトで対応するのも手間がかかる。

調べてみるとTF_FORCE_GPU_ALLOW_GROWTHという環境変数で制御できることがわかり、楽にGPUメモリの制限を扱えるようになった。

$ TF_FORCE_GPU_ALLOW_GROWTH=true python main.py

全ての場合でGPUメモリを制限したいなら、.bash_profileに書き込んでしまってもOK。

export TF_FORCE_GPU_ALLOW_GROWTH=true

参考

Use a GPU | TensorFlow Core