tf.distribute.ReductionToOneDevice
Always do reduction to one device first and then do broadcasting.
Inherits From: CrossDeviceOps
tf.distribute.ReductionToOneDevice(
reduce_to_device=None, accumulation_fn=None
)
Batch reduction is done by reduction on each element one by one.
| 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.
| Returns |
|---|
| a Mirrored object. |
| Raises |
|---|
ValueError | if per_replica_value can't be converted to a PerReplica object. |