[解決済み] TensorflowがGPUのメモリ全体を割り当てるのを防ぐには？

質問

私は、Nvidia Titan X GPUを搭載したサーバーマシンを数台所有し、計算機資源を共有している環境で働いています。

小規模から中規模のモデルであれば、Titan Xの12GBで、通常2-3人が同じGPUで同時にトレーニングを実行するのに十分です。1つのモデルがGPUのすべての計算ユニットをフルに活用できないほど小さい場合、これは実際に、1つのトレーニングプロセスを次々に実行する場合と比較して、高速化につながる可能性があります。GPUへの同時アクセスによって個々の学習時間が遅くなる場合でも、複数のユーザーが同時にGPUで学習できる柔軟性があるのは良いことです。

TensorFlowの問題は、デフォルトでは起動時に利用可能なGPUメモリをフルに割り当ててしまうことです。小さな2層のニューラルネットワークであっても、12GBのGPUメモリをすべて使い切ってしまうんですね。

TensorFlowは、与えられたモデルに対して、例えば4GBのGPUメモリしか割り当てないようにする方法はありますか？

どのように解決するのですか？

を構築する際に、割り当てるGPUメモリの割合を設定することができます。 tf.Session を渡すことで tf.GPUOptions の一部として、オプションの config 引数を指定します。

# Assume that you have 12GB of GPU memory and want to allocate ~4GB:
gpu_options = tf.GPUOptions(per_process_gpu_memory_fraction=0.333)

sess = tf.Session(config=tf.ConfigProto(gpu_options=gpu_options))

は per_process_gpu_memory_fraction は、同じマシン上の各 GPU のプロセスで使用される GPU メモリの量のハード上限として機能します。現在、この割合は同じマシン上のすべての GPU に一様に適用され、GPU ごとに設定する方法はありません。