tf.distribute.Strategy

A state & compute distribution policy on a list of devices.

See the guide for overview and examples. See tf.distribute.StrategyExtended and tf.distribute for a glossory of concepts mentioned on this page such as "per-replica", replica, and reduce.

In short:

A custom training loop can be as simple as:

with my_strategy.scope():
  @tf.function
  def distribute_train_epoch(dataset):
    def replica_fn(input):
      # process input and return result
      return result

    total_result = 0
    for x in dataset:
      per_replica_result = my_strategy.run(replica_fn, args=(x,))
      total_result += my_strategy.reduce(tf.distribute.ReduceOp.SUM,
                                         per_replica_result, axis=None)
    return total_result

  dist_dataset = my_strategy.experimental_distribute_dataset(dataset)
  for _ in range(EPOCHS):
    train_result = distribute_train_epoch(dist_dataset)

This takes an ordinary dataset and replica_fn and runs it distributed using a particular tf.distribute.Strategy named my_strategy above. Any variables created in replica_fn are created using my_strategy's policy, and library functions called by replica_fn can use the get_replica_context() API to implement distributed-specific behavior.

You can use the reduce API to aggregate results across replicas and use this as a return value from one iteration over a tf.distribute.DistributedDataset. Or you can use tf.keras.metrics (such as loss, accuracy, etc.) to accumulate metrics across steps in a given epoch.

See the custom training loop tutorial for a more detailed example.

cluster_resolver Returns the cluster resolver associated with this strategy.

In general, when using a multi-worker tf.distribute strategy such as tf.distribute.experimental.MultiWorkerMirroredStrategy or tf.distribute.experimental.TPUStrategy(), there is a tf.distribute.cluster_resolver.ClusterResolver associated with the strategy used, and such an instance is returned by this property.

Strategies that intend to have an associated tf.distribute.cluster_resolver.ClusterResolver must set the relevant attribute, or override this property; otherwise, None is returned by default. Those strategies should also provide information regarding what is returned by this property.

Single-worker strategies usually do not have a tf.distribute.cluster_resolver.ClusterResolver, and in those cases this property will return None.

The tf.distribute.cluster_resolver.ClusterResolver may be useful when the user needs to access information such as the cluster spec, task type or task id. For example,