tf.distribute.ReductionToOneDevice

View source on GitHub

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

Inherits From: CrossDeviceOps

View aliases

Compat aliases for migration

See Migration guide for more details.

tf.compat.v1.distribute.ReductionToOneDevice

tf.distribute.ReductionToOneDevice(
    reduce_to_device=None, accumulation_fn=None
)

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

mirrored_strategy = tf.distribute.MirroredStrategy(
  cross_device_ops=tf.distribute.ReductionToOneDevice())

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, experimental_hints=None
)

Reduce PerReplica objects in a batch.

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

This can be faster than multiple individual reduces because we can fuse several tensors into one or multiple packs before reduction.

Args
reduce_op	An instance of tf.distribute.ReduceOp that indicates how the per_replica_value will be reduced.
value_destination_pairs	A list or a tuple of PerReplica objects (or tensors with device set if there is one device) and destinations.
experimental_hints	A tf.distrbute.experimental.CollectiveHints. Hints to perform collective operations.

Returns
a list of Mirrored objects.

Raises
ValueError	if value_destination_pairs is not an iterable 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, experimental_hints=None
)

Reduce per_replica_value to destinations.

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

Args
reduce_op	An instance of tf.distribute.ReduceOp that indicates how per_replica_value will be reduced.
per_replica_value	A tf.distribute.DistributedValues object or a tensor with device set.
destinations	the reduction destinations.
experimental_hints	A tf.distrbute.experimental.CollectiveHints. Hints to perform collective operations.

Returns
a Mirrored object.

Raises
ValueError	if per_replica_value can't be converted to a PerReplica object or if destinations aren't strings, Variables or DistributedValues