def _moments(self, inputs, reduction_axes, keep_dims): """Compute the mean and variance: it overrides the original _moments.""" worker_mean, worker_variance = super( SyncBatchNormalization, self)._moments(inputs, reduction_axes, keep_dims=keep_dims) # print(worker_mean) # print(worker_variance) if size() > 1: # Compute variance using: Var[X] = E[X^2] - E[X]^2. worker_square_of_mean = tf.math.square(worker_mean) worker_mean_of_square = worker_variance + worker_square_of_mean # Average stats across all workers worker_stack = tf.stack([worker_mean, worker_mean_of_square]) group_stack = _allreduce(worker_stack, op=Average) group_stack /= size() group_mean, group_mean_of_square = tf.unstack(group_stack) group_variance = group_mean_of_square - tf.math.square(group_mean) return (group_mean, group_variance) else: return (worker_mean, worker_variance)
def allreduce(tensor, average=True, device_dense='', device_sparse='', compression=Compression.none): """Perform an allreduce on a tf.Tensor or tf.IndexedSlices. This function performs a bandwidth-optimal ring allreduce on the input tensor. If the input is an tf.IndexedSlices, the function instead does an allgather on the values and the indices, effectively doing an allreduce on the represented tensor. Arguments: tensor: tf.Tensor, tf.Variable, or tf.IndexedSlices to reduce. The shape of the input must be identical across all ranks. average: If True, computes the average over all ranks. Otherwise, computes the sum over all ranks. device_dense: Device to be used for dense tensors. Uses GPU by default if Horovod was built with HOROVOD_GPU_ALLREDUCE. device_sparse: Device to be used for sparse tensors. Uses GPU by default if Horovod was built with HOROVOD_GPU_ALLGATHER. compression: Compression algorithm used to reduce the amount of data sent and received by each worker node. Defaults to not using compression. Returns: A tensor of the same shape and type as `tensor`, summed across all processes. """ if isinstance(tensor, tf.IndexedSlices): with tf.device(device_sparse): # For IndexedSlices, do two allgathers instead of an allreduce. horovod_size = tf.cast(size(), tensor.values.dtype) values = allgather(tensor.values) indices = allgather(tensor.indices) # To make this operation into an average, divide allgathered values by # the Horovod size. new_values = tf.div(values, horovod_size) if average else values return tf.IndexedSlices(new_values, indices, dense_shape=tensor.dense_shape) else: with tf.device(device_dense): horovod_size = tf.cast(size(), dtype=tensor.dtype) tensor_compressed, ctx = compression.compress(tensor) summed_tensor_compressed = _allreduce(tensor_compressed) summed_tensor = compression.decompress(summed_tensor_compressed, ctx) new_tensor = (tf.div(summed_tensor, horovod_size) if average else summed_tensor) return new_tensor
def allreduce(tensor, average_dense=True, average_sparse=True, use_allgatherv=False, device_dense='', device_sparse=''): """Perform an allreduce on a tf.Tensor or tf.IndexedSlices. Arguments: tensor: tf.Tensor, tf.Variable, or tf.IndexedSlices to reduce. The shape of the input must be identical across all ranks. average_dense: If True, computes the average over all ranks for dense tensors. Otherwise, computes the sum over all ranks. average_sparse: If True, computes the average over all ranks for sparse tensors. Otherwise, computes the sum over all ranks. device_dense: Device to be used for dense tensors. Uses GPU by default if Horovod was build with HOROVOD_GPU_ALLREDUCE. device_sparse: Device to be used for sparse tensors. Uses GPU by default if Horovod was build with HOROVOD_GPU_ALLGATHER. This function performs a bandwidth-optimal ring allreduce on the input tensor. If the input is an tf.IndexedSlices, the function instead does an allgather on the values and the indices, effectively doing an allreduce on the represented tensor. """ if isinstance(tensor, tf.IndexedSlices): with tf.device(device_sparse): # For IndexedSlices, do two allgathers intead of an allreduce. if use_allgatherv: values = allgatherv(tensor.values) indices = allgatherv(tensor.indices) else: values = allgather(tensor.values) indices = allgather(tensor.indices) # To make this operation into an average, divide all gathered values by # the Horovod size. horovod_size = tf.cast(size(), tensor.values.dtype) new_values = tf.div(values, horovod_size) if average_sparse else values return tf.IndexedSlices(new_values, indices, dense_shape=tensor.dense_shape) else: with tf.device(device_dense): horovod_size = tf.cast(size(), tensor.dtype) summed_tensor = _allreduce(tensor) new_tensor = (tf.div(summed_tensor, horovod_size) if average_dense else summed_tensor) return new_tensor
def allreduce(tensor, average=True, device_dense='', device_sparse=''): """Perform an allreduce on a tf.Tensor or tf.IndexedSlices. Arguments: tensor: tf.Tensor, tf.Variable, or tf.IndexedSlices to reduce. The shape of the input must be identical across all ranks. average: If True, computes the average over all ranks. Otherwise, computes the sum over all ranks. device_dense: Device to be used for dense tensors. Uses GPU by default if Horovod was build with HOROVOD_GPU_ALLREDUCE. device_sparse: Device to be used for sparse tensors. Uses GPU by default if Horovod was build with HOROVOD_GPU_ALLGATHER. This function performs a bandwidth-optimal ring allreduce on the input tensor. If the input is an tf.IndexedSlices, the function instead does an allgather on the values and the indices, effectively doing an allreduce on the represented tensor. """ if isinstance(tensor, tf.IndexedSlices): with tf.device(device_sparse): # For IndexedSlices, do two allgathers intead of an allreduce. horovod_size = tf.cast(size(), tensor.values.dtype) values = allgather(tensor.values) indices = allgather(tensor.indices) # To make this operation into an average, divide all gathered values by # the Horovod size. new_values = tf.div(values, horovod_size) if average else values return tf.IndexedSlices(new_values, indices, dense_shape=tensor.dense_shape) else: with tf.device(device_dense): horovod_size = tf.cast(size(), tensor.dtype) summed_tensor = _allreduce(tensor) new_tensor = (tf.div(summed_tensor, horovod_size) if average else summed_tensor) return new_tensor
def allreduce(tensor, average=None, device_dense='', device_sparse='', compression=Compression.none, op=None, prescale_factor=1.0, postscale_factor=1.0, name=None): """Perform an allreduce on a tf.Tensor or tf.IndexedSlices. This function performs a bandwidth-optimal ring allreduce on the input tensor. If the input is an tf.IndexedSlices, the function instead does an allgather on the values and the indices, effectively doing an allreduce on the represented tensor. Arguments: tensor: tf.Tensor, tf.Variable, or tf.IndexedSlices to reduce. The shape of the input must be identical across all ranks. average: .. warning:: .. deprecated:: 0.19.0 Use `op` instead. Will be removed in v0.21.0. device_dense: Device to be used for dense tensors. Uses GPU by default if Horovod was built with HOROVOD_GPU_OPERATIONS. device_sparse: Device to be used for sparse tensors. Uses GPU by default if Horovod was built with HOROVOD_GPU_OPERATIONS. compression: Compression algorithm used to reduce the amount of data sent and received by each worker node. Defaults to not using compression. op: The reduction operation to combine tensors across different ranks. Defaults to Average if None is given. prescale_factor: Multiplicative factor to scale tensor before allreduce. postscale_factor: Multiplicative factor to scale tensor after allreduce. name: A name of the allreduce operation Returns: A tensor of the same shape and type as `tensor`, summed across all processes. """ op = handle_average_backwards_compatibility(op, average) if isinstance(tensor, tf.IndexedSlices): # TODO: Need to fix this to actuall call Adasum if op == Adasum: raise NotImplementedError( 'The Adasum reduction does not currently support sparse tensors. As a ' 'workaround please pass sparse_as_dense=True to DistributedOptimizer' ) with tf.device(device_sparse): # For IndexedSlices, do two allgathers instead of an allreduce. horovod_size = tf.cast(size_op() if int( os.environ.get("HOROVOD_ELASTIC", 0)) else size(), dtype=tensor.values.dtype) values = allgather(tensor.values) indices = allgather(tensor.indices) # To make this operation into an average, divide allgathered values by # the Horovod size. new_values = (values / horovod_size) if op == Average else values return tf.IndexedSlices(new_values, indices, dense_shape=tensor.dense_shape) else: average_in_framework = False if rocm_built(): # For ROCm, perform averaging at framework level average_in_framework = op == Average or op == Adasum op = Sum if op == Average else op with tf.device(device_dense): horovod_size = tf.cast(size_op() if int( os.environ.get("HOROVOD_ELASTIC", 0)) else size(), dtype=tensor.dtype) tensor_compressed, ctx = compression.compress(tensor) summed_tensor_compressed = _allreduce( tensor_compressed, op=op, prescale_factor=prescale_factor, postscale_factor=postscale_factor, name=name) summed_tensor = compression.decompress(summed_tensor_compressed, ctx) if op == Adasum: if 'CPU' not in tensor.device and gpu_available('tensorflow'): if nccl_built(): if not is_homogeneous: raise NotImplementedError( 'Running GPU Adasum on heterogeneous cluster is not supported yet.' ) elif not check_num_rank_power_of_2( int(size() / local_size())): raise NotImplementedError( 'Running GPU Adasum with non-power of 2 nodes is not supported yet.' ) if rocm_built(): horovod_local_size = tf.cast( local_size_op() if int( os.environ.get("HOROVOD_ELASTIC", 0)) else local_size(), dtype=tensor.dtype) new_tensor = summed_tensor / horovod_local_size else: new_tensor = summed_tensor else: warnings.warn( 'Adasum reduction does not currently support GPU reduction using MPI. Tensors ' 'are copied to CPU memory instead. To use Adasum for GPU reduction, please ' 'compile Horovod with HOROVOD_GPU_OPERATIONS=NCCL.' ) new_tensor = summed_tensor else: if not check_num_rank_power_of_2(size()): raise NotImplementedError( 'Running Adasum with non-power of 2 ranks is not supported yet.' ) new_tensor = summed_tensor else: if rocm_built(): new_tensor = (summed_tensor / horovod_size ) if average_in_framework else summed_tensor else: new_tensor = summed_tensor return new_tensor
def allreduce(tensor, average=None, device_dense='', device_sparse='', compression=Compression.none, op=None): """Perform an allreduce on a tf.Tensor or tf.IndexedSlices. This function performs a bandwidth-optimal ring allreduce on the input tensor. If the input is an tf.IndexedSlices, the function instead does an allgather on the values and the indices, effectively doing an allreduce on the represented tensor. Arguments: tensor: tf.Tensor, tf.Variable, or tf.IndexedSlices to reduce. The shape of the input must be identical across all ranks. average: DEPRECATED, please use op instead. device_dense: Device to be used for dense tensors. Uses GPU by default if Horovod was built with HOROVOD_GPU_ALLREDUCE. device_sparse: Device to be used for sparse tensors. Uses GPU by default if Horovod was built with HOROVOD_GPU_ALLGATHER. compression: Compression algorithm used to reduce the amount of data sent and received by each worker node. Defaults to not using compression. op: The reduction operation to combine tensors across different ranks. Defaults to Average if None is given. Returns: A tensor of the same shape and type as `tensor`, summed across all processes. """ op = handle_average_backwards_compatibility(op, average) # Averaging happens in framework code, so translate that to Sum for the actual call true_op = Sum if op == Average else op if isinstance(tensor, tf.IndexedSlices): # TODO: Need to fix this to actuall call Adasum if op == Adasum: raise NotImplementedError( "The Adasum reduction does not currently support " "sparse tensors. As a workaround please pass sparse_as_dense=True to " "DistributedOptimizer") with tf.device(device_sparse): # For IndexedSlices, do two allgathers instead of an allreduce. horovod_size = tf.cast(size(), tensor.values.dtype) values = allgather(tensor.values) indices = allgather(tensor.indices) # To make this operation into an average, divide allgathered values by # the Horovod size. new_values = (values / horovod_size) if op == Average else values return tf.IndexedSlices(new_values, indices, dense_shape=tensor.dense_shape) else: with tf.device(device_dense): horovod_size = tf.cast(size(), dtype=tensor.dtype) tensor_compressed, ctx = compression.compress(tensor) summed_tensor_compressed = _allreduce(tensor_compressed, op=true_op) summed_tensor = compression.decompress(summed_tensor_compressed, ctx) if op == Adasum: if ('CPU' not in tensor.device and has_gpu): if nccl_built(): if not is_homogeneous: raise NotImplementedError( 'Running GPU Adasum on heterogeneous cluster is not supported yet.' ) elif not check_num_rank_power_of_2( int(size() / local_size())): raise NotImplementedError( 'Running GPU Adasum with non-power of 2 nodes is not supported yet.' ) horovod_local_size = tf.cast(local_size(), dtype=tensor.dtype) new_tensor = summed_tensor / horovod_local_size else: warnings.warn( "Adasum reduction does not currently support " "GPU reduction using MPI. Tensors are copied to CPU memory instead." "To use Adasum for GPU reduction, please compile Horovod with HOROVOD_GPU_ALLREDUCE=NCCL." ) new_tensor = summed_tensor else: if not check_num_rank_power_of_2(size()): raise NotImplementedError( 'Running Adasum with non-power of 2 ranks is not supported yet.' ) new_tensor = summed_tensor else: new_tensor = (summed_tensor / horovod_size) if op == Average else summed_tensor return new_tensor