TensorFlow 1 version View source on GitHub

A one-machine strategy that puts all variables on a single device.

Inherits From: Strategy

Used in the notebooks

Used in the guide

Variables are assigned to local CPU or the only GPU. If there is more than one GPU, compute operations (other than variable update operations) will be replicated across all GPUs.

For Example:

strategy = tf.distribute.experimental.CentralStorageStrategy()
# Create a dataset
ds =
# Distribute that dataset
dist_dataset = strategy.experimental_distribute_dataset(ds)

with strategy.scope():
  def train_step(val):
    return val + 1

  # Iterate over the distributed dataset
  for x in dist_dataset:
    # process dataset elements, args=(x,))

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,

os.environ['TF_CONFIG'] = json.dumps({
'cluster': {
'worker': ["localhost:12345", "localhost:23456"],
'ps': ["localhost:34567"]
'task': {'type': 'worker', 'index': 0}

# This implicitly uses TF_CONFIG for the cluster and current task info.
strategy = tf.distribute.experimental.MultiWorkerMirroredStrategy()


if strategy.cluster_resolver.task_type == 'worker':
# Perform something that's only applicable on workers. Since we set this
# as a worker above, this block will run on this particular instance.
elif strategy.cluster_resolver.task_type == 'ps':
# Perform something that's only applicable on parameter servers. Since we
# set this as a worker above, this block will not run on this particular
# instance.

For more information, please see tf.distribute.cluster_resolver.ClusterResolver's API docstring.

extended tf.distribute.StrategyExtended with additional methods.
num_replicas_in_sync Returns number of replicas over which gradients are aggregated.



View source

Adds annotation that tensor will be assigned to a logical device.

# Initializing TPU system with 2 logical devices and 4 replicas.
resolver = tf.distribute.cluster_resolver.TPUClusterResolver(tpu='')
topology = tf.tpu.experimental.initialize_tpu_system(resolver)
device_assignment =
    computation_shape=[1, 1, 1, 2],
strategy = tf.distribute.TPUStrategy(
    resolver, experimental_device_assignment=device_assignment)
iterator = iter(inputs)

def step_fn(inputs):
  output = tf.add(inputs, inputs)

  # Add operation will be executed on logical device 0.
  output = strategy.experimental_assign_to_logical_device(output, 0)
  return output, args=(next(iterator),))

tensor Input tensor to annotate.
logical_device_id Id of the logical core to which the tensor will be assigned.

ValueError The logical device id presented is not consistent with total number of partitions specified by the device assignment.

Annotated tensor with idential value as tensor.


View source

Distributes a instance provided via dataset.

The returned dataset is a wrapped strategy dataset which creates a multidevice iterator under the hood. It prefetches the input data to the specified devices on the worker. The returned distributed dataset can be iterated over similar to how regular datasets can.

For Example:

strategy = tf.distribute.CentralStorageStrategy()  # with 1 CPU and 1 GPU
dataset =
dist_dataset = strategy.experimental_distribute_dataset(dataset)
for x in dist_dataset:
  print(x)  # Prints PerReplica values [0, 1], [2, 3],...

Args: dataset: to be prefetched to device.

A "distributed Dataset" that the caller can iterate over.


View source

Distributes instances created by calls to dataset_fn.

dataset_fn will be called once for each worker in the strategy. In this case, we only have one worker so dataset_fn is called once. Each replica on this worker will then dequeue a batch of elements from this local dataset.

The dataset_fn should take an tf.distribute.InputContext instance where information about batching and input replication can be accessed.

For Example:

def dataset_fn(input_context):
  batch_size = input_context.get_per_replica_batch_size(global_batch_size)
  d =[[1.]]).repeat().batch(batch_size)
  return d.shard(
      input_context.num_input_pipelines, input_context.input_pipeline_id)

inputs = strategy.experimental_distribute_datasets_from_function(dataset_fn)

for batch in inputs:
  replica_results =, args=(batch,))

dataset_fn A function taking a tf.distribute.InputContext instance and returning a

A "distributed Dataset", which the caller can iterate over like regular datasets.


View source

Generates tf.distribute.DistributedValues from value_fn.

This function is to generate tf.distribute.DistributedValues to pass into run, reduce, or other methods that take distributed values when not using datasets.

value_fn The function to run to generate values. It is called for each replica with tf.distribute.ValueContext as the sole argument. It must return a Tensor or a type that can be converted to a Tensor.

A tf.distribute.DistributedValues containing a value for each replica.

Example usage:

  1. Return constant value per replica:
strategy = tf.distribute.MirroredStrategy()
def value_fn(ctx):
  return tf.constant(1.)
distributed_values = (
local_result = strategy.experimental_local_results(distributed_values)
(<tf.Tensor: shape=(), dtype=float32, numpy=1.0>,)
  1. Distribute values in array based on replica_id:
strategy = tf.distribute.MirroredStrategy()
array_value = np.array([3., 2., 1.])
def value_fn(ctx):
  return array_value[ctx.replica_id_in_sync_group]
distributed_values = (
local_result = strategy.experimental_local_results(distributed_values)
  1. Specify values using num_replicas_in_sync:
strategy = tf.distribute.MirroredStrategy()
def value_fn(ctx):
  return ctx.num_replicas_in_sync
distributed_values = (
local_result = strategy.experimental_local_results(distributed_values)
  1. Place values on devices and distribute:
strategy = tf.distribute.TPUStrategy()
worker_devices = strategy.extended.worker_devices
multiple_values = []
for i in range(strategy.num_replicas_in_sync):
  with tf.device(worker_devices[i]):

def value_fn(ctx):
  return multiple_values[ctx.replica_id_in_sync_group]

distributed_values = strategy.


View source

Returns the list of all local per-replica values contained in value.

In CentralStorageStrategy there is a single worker so the value returned will be all the values on that worker.

value A value returned by run(), extended.call_for_each_replica(), or a variable created in scope.

A tuple of values contained in value. If value represents a single value, this returns (value,).


View source

Makes a from a numpy array. (deprecated)

This avoids adding numpy_input as a large constant in the graph, and copies the data to the machine or machines that will be processing the input.

Note that you will likely need to use experimental_distribute_dataset with the returned dataset to further distribute it with the strategy.


strategy = tf.distribute.MirroredStrategy()