tf.distribute.ReplicaContext

View source on GitHub

tf.distribute.Strategy API when in a replica context.

tf.distribute.ReplicaContext(
    strategy, replica_id_in_sync_group
)

You can use tf.distribute.get_replica_context to get an instance of ReplicaContext. This should be inside your replicated step function, such as in a tf.distribute.Strategy.experimental_run_v2 call.

Attributes:

Methods

__enter__

View source

__enter__()

__exit__

View source

__exit__(
    exception_type, exception_value, traceback
)

all_reduce

View source

all_reduce(
    reduce_op, value
)

All-reduces the given value Tensor nest across replicas.

If all_reduce is called in any replica, it must be called in all replicas. The nested structure and Tensor shapes must be identical in all replicas.

IMPORTANT: The ordering of communications must be identical in all replicas.

Example with two replicas: Replica 0 value: {'a': 1, 'b': [40, 1]} Replica 1 value: {'a': 3, 'b': [ 2, 98]}

If reduce_op == SUM: Result (on all replicas): {'a': 4, 'b': [42, 99]}

If reduce_op == MEAN: Result (on all replicas): {'a': 2, 'b': [21, 49.5]}

Args:

Returns:

A Tensor nest with the reduced values from each replica.

merge_call

View source

merge_call(
    merge_fn, args=(), kwargs=None
)

Merge args across replicas and run merge_fn in a cross-replica context.

This allows communication and coordination when there are multiple calls to the step_fn triggered by a call to strategy.experimental_run_v2(step_fn, ...).

See tf.distribute.Strategy.experimental_run_v2 for an explanation.

If not inside a distributed scope, this is equivalent to:

strategy = tf.distribute.get_strategy()
with cross-replica-context(strategy):
  return merge_fn(strategy, *args, **kwargs)

Args:

Returns:

The return value of merge_fn, except for PerReplica values which are unpacked.