Tensorflow what is the tf.contrib.nccl.allsum in new version?

user_9881203 0
Viewed 773

It seems that from tensorflow 1.13, there is no api such as tf.contrib.nccl.allsum. However, in the Nvidia official GitHub https://github.com/tkarras/progressive_growing_of_gans, which uses this old API to reduce sum from different gpu devices as the following. 

# 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])

I am not sure if there is similar api which can achieve the same collective operation cross different devices. I have checked the Tensorflow official website and it seems that programmers prefer to use tf.distribute.MirroredStrategy which hides the raw operation of NCCL. Thanks a lot.

tensorflow gpu
1 Answers

I think the same API is nccl_ops.all_sum. I have demoed this API by the following code. 

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))

I am not sure what tensorflow team will do in the future. But now we can use it to do the collective operations.

user_98812030
Related