我在共享計算資源的環境中工作,也就是說,我們有幾台服務器機器都配備了幾個Nvidia Titan X GPU。
對於小到中等尺寸的模型,Titan X的12GB通常足以讓2-3人在同一GPU上同時進行訓練。如果模型足夠小,以至於單個模型不能充分利用Titan X的所有計算單元,那麽實際上可以導致某種加速(相比串行多個訓練任務來說)。即使在GPU的並發訪問確實減慢了單個訓練時間的情況下,同時在GPU上擁有多個用戶運行的靈活性仍然不錯。
TensorFlow的問題在於,默認情況下,它會在GPU啟動時為其分配全部可用內存。即使對於一個小型的2層神經網絡,我也看到Titan X的12 GB用完了。
有沒有辦法讓TensorFlow隻分配4GB的GPU內存,如果知道這個數量對於給定的模型是足夠的?
最佳解決辦法
通過將tf.GPUOptions
作為可選config
參數的一部分傳遞來構造tf.Session
時,可以設置要分配的GPU內存部分:
# 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上;沒有辦法在per-GPU的基礎上進行設置。
次佳解決辦法
config = tf.ConfigProto()
config.gpu_options.allow_growth=True
sess = tf.Session(config=config)
https://github.com/tensorflow/tensorflow/issues/1578