python - TensorFlow 分布式运行时模型并行 CIFAR-10

Question

我已尝试修改 CIFAR-10 示例以在新的 TensorFlow 分布式运行时上运行。但是，在尝试运行该程序时出现以下错误:

InvalidArgumentError: Cannot assign a device to node 'softmax_linear/biases/ExponentialMovingAverage': 
Could not satisfy explicit device specification '/job:local/task:0/device:CPU:0'

我使用以下命令启动集群。在我运行的第一个节点上:

bazel-bin/tensorflow/core/distributed_runtime/rpc/grpc_tensorflow_server --cluster_spec='local|10.31.101.101:7777;10.31.101.224:7778' --job_name=local --task_id=0

...在我运行的第二个节点上:

bazel-bin/tensorflow/core/distributed_runtime/rpc/grpc_tensorflow_server --cluster_spec='local|10.31.101.101:7777;10.31.101.224:7778' --job_name=local --task_id=1

对于CIFAR-10 multi-GPU code ，我进行了简单的修改，替换了 train() 函数中的两行。下面一行:

with tf.Graph().as_default(), tf.device('/cpu:0'):

...替换为:

with tf.Graph().as_default(), tf.device('/job:local/task:0/cpu:0'):

和下一行:

with tf.device('/gpu:%d' % i):

...替换为:

with tf.device('/job:local/task:0/gpu:%d' % i):

以我的理解，第二个替换应该负责模型替换。运行一个更简单的示例，例如下面的代码，效果很好:

with tf.device('/job:local/task:0/cpu:0'):
    c = tf.constant("Hello, distributed TensorFlow!")
    sess.run(c)
    print(c)

Answer 1

我无法从你的程序中看出，但我的猜测是你还必须修改 line that creates the session指定您的一项工作任务的地址。例如，给定上面的配置，您可以这样写:

sess = tf.Session(
    "grpc://10.31.101.101:7777",
    config=tf.ConfigProto(
        allow_soft_placement=True,
        log_device_placement=FLAGS.log_device_placement))

碰巧的是，我们一直在努力改进该错误消息以减少混淆。如果您在 GitHub 中更新到最新版本并运行相同的代码，您应该会看到一条错误消息，解释为什么无法满足设备规范。