Tensorflow-如何减少/广播集群中具有多个工作线程的张量?

时间:2019-08-01 19:19:14

标签: python tensorflow multi-gpu

我只是想以同步的方式求和张量流中跨工作者的张量。这样的感觉应该非常容易,但是我还没有找到方法。

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

以下是我尝试使用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:我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:集体操作需要设备分配”

1 个答案:

答案 0 :(得分:0)

问题在github上提出,并收到团队成员的答复。有一个代码段可解决此图形间通信问题。