Tensorflow Что такое tf.contrib.nccl.allsum в новой версии?

Question

Похоже, что в тензорном потоке 1.13 нет API, такого как tf.contrib.nccl.allsum. Тем не менее, в официальном 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))

Я не уверен, что команда tenorflow сделает в будущем. Но теперь мы можем использовать его для выполнения коллективных операций.