Tensorflow Module (API Reference)

Warning

The tensorflow module is not finished yet. There is no guarantee for the tensorflow API. Please consider to use pytorch api first.

bluefog.tensorflow.DistributedGradientTape(gradtape, device='')

A tape that wraps another tf.GradientTape, using an allreduce to average gradient values before applying gradients to model weights.

Parameters:
  • gradtape – GradientTape to use for computing gradients and applying updates.
  • device – Device to be used for dense tensors. Uses GPU by default if CUDA is available
bluefog.tensorflow.DistributedOptimizer(optimizer, name=None, use_locking=False, device='')

Construct a new DistributedOptimizer, which uses another optimizer under the hood for computing single-process gradient values and applying gradient updates after the gradient values have been averaged across all the Bluefog ranks.

Parameters:
  • optimizer – Optimizer to use for computing gradients and applying updates.
  • name – Optional name prefix for the operations created when applying gradients. Defaults to “Distributed” followed by the provided optimizer type.
  • use_locking – Whether to use locking when updating variables. See Optimizer.__init__ for more info.
  • device – Device to be used for dense tensors. Uses GPU by default if CUDA is available
bluefog.tensorflow.allgather(tensor: <MagicMock id='140044196795560'>, name: str = None) → <MagicMock id='140044196244952'>

An op which concatenates the input tensor with the same input tensor on all other Bluefog processes.

The concatenation is done on the first dimension, so the input tensors on the different processes must have the same rank and shape, except for the first dimension, which is allowed to be different.

Parameters:
  • tensor – A tensor to allgather.
  • name – A name of the allgather operation.
Returns:

A tensor of the same type as tensor, concatenated on dimension zero across all processes. The shape is identical to the input shape, except for the first dimension, which may be greater and is the sum of all first dimensions of the tensors in different Bluefog processes.
bluefog.tensorflow.allreduce(tensor: <MagicMock id='140044196437408'>, average: bool = True, name: str = None, device: str = '') → <MagicMock id='140044196377880'>

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.

Parameters:
  • 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.
  • name – A name of the allreduce operation.
  • device – Device to be used for dense tensors.
Returns:

A tensor of the same shape and type as tensor, summed across all processes.
bluefog.tensorflow.broadcast(tensor: <MagicMock id='140044196820416'>, root_rank: int, name: str = None) → <MagicMock id='140044196854304'>

An op which broadcasts the input tensor on root rank to the same input tensor on all other Bluefog processes.

The broadcast operation is keyed by the name of the op. The tensor type and shape must be the same on all Bluefog processes for a given name. The broadcast will not start until all processes are ready to send and receive the tensor.

Parameters:
  • tensor – A tensor to broadcast.
  • root_rank – The rank to broadcast the value from.
  • name – A name of the broadcast operation.
Returns:

A tensor of the same shape and type as tensor, with the value broadcasted from root rank.
bluefog.tensorflow.broadcast_variables(variables, root_rank)

Broadcasts variables from root rank to all other processes.

Parameters:
  • variables – variables for broadcast
  • root_rank – rank of the process from which global variables will be broadcasted to all other processes.
bluefog.tensorflow.tf