Tensorflow - 如何在集群中减少/广播具有多个工作人员的张量？

Question

我只是想以同步的方式在 tensorflow 中对跨工作人员的张量求和。感觉这样应该是极其简单的，但是我还没有找到方法。

我的结论是 tf.distribute 中的任何内容都太抽象了，我需要一些较低级别的操作，因为我只想在我的程序中对单个张量进行操作。如果 tf.distribute 中有任何内容可以让我减少/广播每个工人都有一个实例的张量，请纠正我。

这是我尝试使用collective_ops.all_reduce的一些代码

import sys
import tensorflow as tf
from tensorflow.python.ops import collective_ops

task_id = int(sys.argv[1])
cluster = tf.train.ClusterSpec({"worker": ["localhost:2222", "localhost:2223"]})
server = tf.train.Server(cluster, job_name="worker", task_index=task_id)

with tf.device("/job:worker/task:{}/device:CPU:0".format(task_id)):
    t = tf.constant([-1,-3], name='testtensor')
    res = collective_ops.all_reduce(t, group_size=2, group_key=123, instance_key=123, merge_op='Add', final_op='Id', subdiv_offsets=(0,))

with tf.Session(server.target) as sess:

    print('running reduce..')
    print(sess.run(res))

如果我在一个终端中运行上述脚本： python myscript.py 0 和另一个终端中的另一个实例： python myscript.py 1

我希望他们俩都打印总和 [-2, -6]，但他们被困在阻塞调用 less.run(res)

当我只启动了两个进程之一时，它一直在打印“2019-08-01 12:05:24.324155: I tensorflow/core/distributed_runtime/master.cc:268] CreateSession 仍在等待工作人员的响应: /job:worker/replica:0/task:0", 即等待其他工人的回应。 当我开始第二个进程时，上面的日志记录停止了，但无论如何都没有发生任何事情。

我还尝试使用以下方法将张量分配给本地设备 with tf.device(tf.train.replica_device_setter(worker_device="/job:worker/task:0/device:CPU:0", cluster=cluster)): 但是后来我在尝试定义张量时遇到了这个错误： “ValueError：集体操作所需的设备分配”

Answer 1

在github 上提出了一个问题，并收到了团队成员的回复。有一个代码 sn-p 可以解决这个图间通信问题。