Tensorflow contrib.learn.Estimator multi-GPU

Question

爲了使用contrib.learn.Estimator進行多GPU訓練，我試圖在我的model_fn中指定GPU分配。

在僞代碼：

def model_fn(X, y): 
    with tf.device('/gpu:1'): 
     ... various tensorflow ops for model ... 

     return predictions, loss, train_op 

一切正常，沒有tf.device（ '/ GPU：1'）的電話，但它我會遇到以下錯誤：

InvalidArgumentError (see above for traceback): Cannot assign a device to 
node 'save/ShardedFilename_1': Could not satisfy explicit device 
specification '/device:GPU:1' because no supported kernel 
for GPU devices is available. 

我不認爲我自己將圖形添加到圖形，而是通過Estimator的快照功能注入圖形。

我認爲解決方案是設置allow_soft_placement = True，以便非GPU功能將落入CPU，但對於我在處理contrib.learn.Estimator時如何暴露並不明顯。

我看到該選項通常在傳遞給會話的ConfigProto &中設置，但我一直在使用Estimator的功能來爲我管理會話。我是否應該控制會話創建，還是在某個地方缺少一個參數來完成此操作？

非常感謝您的任何建議。

Answer 1

隨着估算器在Tensorflow 1.0中的離開貢獻，這是固定的。