Tensorflow contrib.learn.Estimator 多 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，但我不清楚在处理 [=27= 时它是如何暴露的].

我看到该选项通常在 ConfigProto 中设置并传递给会话，但我一直在使用 Estimator 的功能为我管理会话。我应该控制会话创建，还是我在某处缺少参数来完成此操作？

非常感谢您的任何建议。

Answer 1

随着 Estimator 在 Tensorflow 1.0 中离开 contrib，这已得到修复。