def allreduce_grads(all_grads, average):
    """
    All-reduce average the gradients among K devices. Results are broadcasted to all devices.

    Args:
        all_grads (K x N): List of list of gradients. N is the number of variables.
        average (bool): average gradients or not.

    Returns:
        K x N: same as input, but each grad is replaced by the average over K devices.
    """

    if get_tf_version_tuple() <= (1, 12):
        from tensorflow.contrib import nccl
    else:
        from tensorflow.python.ops import nccl_ops as nccl
    nr_tower = len(all_grads)
    if nr_tower == 1:
        return all_grads
    new_all_grads = []  # N x K
    for grads in zip(*all_grads):
        summed = nccl.all_sum(grads)

        grads_for_devices = []  # K
        for g in summed:
            with tf.device(g.device):
                # tensorflow/benchmarks didn't average gradients
                if average:
                    g = tf.multiply(g, 1.0 / nr_tower)
            grads_for_devices.append(g)
        new_all_grads.append(grads_for_devices)

    # transpose to K x N
    ret = list(zip(*new_all_grads))
    return ret
def allreduce_grads ( all_grads , average ) : if get_tf_version_tuple ( ) <= ( 1 , 12 ) : from tensorflow . contrib import nccl else : from tensorflow . python . ops import nccl_ops as nccl nr_tower = len ( all_grads ) if nr_tower == 1 : return all_grads new_all_grads = [ ] for grads in zip ( * all_grads ) : summed = nccl . all_sum ( grads ) grads_for_devices = [ ] for g in summed : with tf . device ( g . device ) : if average : g = tf . multiply ( g , 1.0 / nr_tower ) grads_for_devices . append ( g ) new_all_grads . append ( grads_for_devices ) ret = list ( zip ( * new_all_grads ) ) return ret
All-reduce average the gradients among K devices. Results are broadcasted to all devices.

    Args:
        all_grads (K x N): List of list of gradients. N is the number of variables.
        average (bool): average gradients or not.

    Returns:
        K x N: same as input, but each grad is replaced by the average over K devices.