新版本中的tf.contrib.nccl.allsum是什么？

Question

从tensorflow 1.13来看，似乎没有像tf.contrib.nccl.allsum这样的api。然而，在Nvidia官方的GitHub https://github.com/tkarras/progressive_growing_of_gans中，它使用这个旧的API来减少来自不同gpu设备的和，如下所示。

# Sum gradients across devices.
            if len(devices) > 1:
                with tf.name_scope('SumAcrossGPUs'), tf.device(None):
                    for var_idx, grad_shape in enumerate(self._grad_shapes):
                        g = [dev_grads[dev][var_idx][0] for dev in devices]
                        if np.prod(grad_shape): # nccl does not support zero-sized tensors
                            g = tf.contrib.nccl.all_sum(g)
                        for dev, gg in zip(devices, g):
                            dev_grads[dev][var_idx] = (gg, dev_grads[dev][var_idx][1])

我不确定是否有类似的api可以实现相同的集体操作，跨越不同的设备。我查看了Tensorflow官方网站，似乎程序员更喜欢使用tf.distribute.MirroredStrategy，它隐藏了NCCL的原始操作。非常感谢。

Answer 1

我认为同样的API是nccl_ops.all_sum。我用下面的代码演示了这个API。

import tensorflow as tf 
from tensorflow.python.ops import nccl_ops

a = []

with tf.device("/GPU:0"):
    g = tf.constant([1,1])
    print(g)
    a.append(g)
with tf.device("/GPU:1"):
    g = tf.constant([2,2])
    a.append(g)

b = nccl_ops.all_sum(a)
with tf.Session() as sess:
    print(sess.run(b))

我不知道tensorflow球队将来会做什么。但现在我们可以用它来做集体行动了。