Missed TensorFlow World? Check out the recap. Learn more

tf.distribute.ReductionToOneDevice

TensorFlow 2.0 version View source on GitHub

Class ReductionToOneDevice

Always do reduction to one device first and then do broadcasting.

Inherits From: CrossDeviceOps

Aliases:

  • Class tf.compat.v1.distribute.ReductionToOneDevice
  • Class tf.compat.v2.distribute.ReductionToOneDevice

Batch reduction is done by reduction on each element one by one.

__init__

View source

__init__(
    reduce_to_device=None,
    accumulation_fn=None
)

Constructor.

Args:

  • reduce_to_device: the intermediate device to reduce to. If None, reduce to the first device in destinations of the reduce() method.
  • accumulation_fn: a function that does accumulation. If None, then tf.math.add_n is used.

Methods

batch_reduce

View source

batch_reduce(
    reduce_op,
    value_destination_pairs
)

Reduce PerReplica objects in a batch.

Reduce each first element in value_destination_pairs to each second element which indicates the destinations.

Args:

  • reduce_op: Indicates how per_replica_value will be reduced. Accepted values are tf.distribute.ReduceOp.SUM, tf.distribute.ReduceOp.MEAN.
  • value_destination_pairs: a list or a tuple of tuples of PerReplica objects (or tensors with device set if there is one device) and destinations.

Returns:

a list of Mirrored objects.

Raises:

  • ValueError: if value_destination_pairs is not a list or a tuple of tuples of PerReplica objects and destinations

broadcast

View source

broadcast(
    tensor,
    destinations
)

Broadcast the tensor to destinations.

Args:

  • tensor: the tensor to broadcast.
  • destinations: the broadcast destinations.

Returns:

a Mirrored object.

reduce

View source

reduce(
    reduce_op,
    per_replica_value,
    destinations
)

Reduce per_replica_value to destinations.

It runs the reduction operation defined by reduce_op and put the result on destinations.

Args:

Returns:

a Mirrored object.

Raises:

  • ValueError: if per_replica_value can't be converted to a PerReplica object.