时间戳

首页 > 解决方案 > Tensorflow 新版本中的 tf.contrib.nccl.allsum 是什么？

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

问题描述

好像从tensorflow 1.13开始，就没有tf.contrib.nccl.allsum之类的api了。然而，在 Nvidia 官方 GitHub https://github.com/tkarras/progressive_growth_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. 非常感谢。

标签： tensorflowgpu

解决方案

我认为相同的 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 团队将来会做什么。但是现在我们可以用它来做集体操作。

推荐阅读