TensorFlow新版本中tf.contrib.nccl.all_sum替代方案咨询

你提到tf.contrib.nccl.all_sum在TensorFlow 1.13之后被移除，又不想用封装性过强的tf.distribute.MirroredStrategy，想要更直接的替代方案，这里有几个适配性拉满的选项：

1. 用tf.nccl.all_sum直接替换（最优方案） #

TensorFlow后续把原contrib下的NCCL模块迁移到了核心API中，tf.nccl.all_sum就是旧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.nccl.all_sum(g)  # 替换为核心API中的NCCL求和
            for dev, gg in zip(devices, g):
                dev_grads[dev][var_idx] = (gg, dev_grads[dev][var_idx][1])

除了all_sum，tf.nccl还提供all_prod、all_max等常用集合操作，和旧版API用法完全一致，几乎零迁移成本。如果是在TensorFlow 1.x高版本（1.13-1.15），可以用tf.compat.v1.nccl.all_sum来保证兼容性。

2. 手动实现跨设备求和（无NCCL依赖场景） #

如果你的环境没法使用NCCL，也可以手动通过tf.reduce_sum结合张量复制实现，虽然性能不如NCCL硬件优化，但逻辑简单可靠：

# 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):
                # 手动完成求和：先聚合所有梯度到同一设备，求和后再分发回各个GPU
                summed_grad = tf.reduce_sum(g, axis=0)
                g = [tf.identity(summed_grad) for _ in devices]
            for dev, gg in zip(devices, g):
                dev_grads[dev][var_idx] = (gg, dev_grads[dev][var_idx][1])

内容的提问来源于stack exchange，提问作者Sean