Tensorflow が GPU メモリ全体を割り当てるのを防ぐにはどうすればよいでしょうか? 質問する

Tensorflow が GPU メモリ全体を割り当てるのを防ぐにはどうすればよいでしょうか? 質問する

私は計算リソースを共有する環境で働いています。つまり、それぞれに数個の Nvidia Titan X GPU を搭載したサーバー マシンが数台あります。

小規模から中規模のモデルの場合、Titan X の 12 GB は通常、2 ~ 3 人が同じ GPU で同時にトレーニングを実行するのに十分です。モデルが小さく、1 つのモデルで GPU のすべての計算ユニットを最大限に活用できない場合は、トレーニング プロセスを 1 つずつ実行するよりも実際にスピードアップできます。GPU への同時アクセスによって個々のトレーニング時間が遅くなる場合でも、複数のユーザーが同時に GPU でトレーニングできる柔軟性は依然として便利です。

TensorFlow の問題は、デフォルトでは起動時に利用可能な GPU メモリの全量が割り当てられることです。小さな 2 層ニューラル ネットワークの場合でも、12 GB の GPU メモリがすべて使用されていることがわかります。

特定のモデルにこれで十分であることがわかっている場合、TensorFlow に、たとえば 4 GB の GPU メモリのみを割り当てるようにする方法はありますか?

ベストアンサー1

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 ごとに設定する方法はありません。

おすすめ記事