tf.distribute.OneDeviceStrategy

A distribution strategy for running on a single device.

Inherits From: Strategy

tf.distribute.OneDeviceStrategy(
    device
)

Using this strategy will place any variables created in its scope on the specified device. Input distributed through this strategy will be prefetched to the specified device. Moreover, any functions called via strategy.run will also be placed on the specified device as well.

Typical usage of this strategy could be testing your code with the tf.distribute.Strategy API before switching to other strategies which actually distribute to multiple devices/machines.

For example:

strategy = tf.distribute.OneDeviceStrategy(device="/gpu:0")

with strategy.scope():
  v = tf.Variable(1.0)
  print(v.device)  # /job:localhost/replica:0/task:0/device:GPU:0

def step_fn(x):
  return x * 2

result = 0
for i in range(10):
  result += strategy.run(step_fn, args=(i,))
print(result)  # 90

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.

Methods

experimental_assign_to_logical_device

View source

experimental_assign_to_logical_device(
    tensor, logical_device_id
)

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

Note: This API is only supported in TPUStrategy for now. This adds an annotation to tensor specifying that operations on tensor will be invoked on logical core device id logical_device_id. When model parallelism is used, the default behavior is that all ops are placed on zero-th logical device.

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

@tf.function()
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

strategy.run(step_fn, args=(next(iterator),))

experimental_distribute_dataset

View source

experimental_distribute_dataset(
    dataset
)

Distributes a tf.data.Dataset instance provided via dataset.

In this case, there is only one device, so this is only a thin wrapper around the input dataset. It will, however, prefetch the input data to the specified device. The returned distributed dataset can be iterated over similar to how regular datasets can.

Note: Currently, the user cannot add any more transformations to a distributed dataset.

Example:

strategy = tf.distribute.OneDeviceStrategy()
dataset = tf.data.Dataset.range(10).batch(2)
dist_dataset = strategy.experimental_distribute_dataset(dataset)
for x in dist_dataset:
  print(x)  # [0, 1], [2, 3],...

Args: dataset: tf.data.Dataset to be prefetched to device.

experimental_distribute_datasets_from_function

View source

experimental_distribute_datasets_from_function(
    dataset_fn
)

Distributes tf.data.Dataset 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 and one device so dataset_fn is called once.

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

def dataset_fn(input_context):
  batch_size = input_context.get_per_replica_batch_size(global_batch_size)
  d = tf.data.Dataset.from_tensors([[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 = strategy.run(replica_fn, args=(batch,))

experimental_distribute_values_from_function

View source

experimental_distribute_values_from_function(
    value_fn
)

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.

Example usage:

1. Return constant value per replica:

strategy = tf.distribute.MirroredStrategy()
def value_fn(ctx):
  return tf.constant(1.)
distributed_values = (
     strategy.experimental_distribute_values_from_function(
       value_fn))
local_result = strategy.experimental_local_results(distributed_values)
local_result
(<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 = (
     strategy.experimental_distribute_values_from_function(
       value_fn))
local_result = strategy.experimental_local_results(distributed_values)
local_result
(3.0,)

1. Specify values using num_replicas_in_sync:

strategy = tf.distribute.MirroredStrategy()
def value_fn(ctx):
  return ctx.num_replicas_in_sync
distributed_values = (
     strategy.experimental_distribute_values_from_function(
       value_fn))
local_result = strategy.experimental_local_results(distributed_values)
local_result
(1,)

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]):
    multiple_values.append(tf.constant(1.0))

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

distributed_values = strategy.
  experimental_distribute_values_from_function(
  value_fn)

experimental_local_results

View source

experimental_local_results(
    value
)

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

In OneDeviceStrategy, the value is always expected to be a single value, so the result is just the value in a tuple.

experimental_make_numpy_dataset

View source

experimental_make_numpy_dataset(
    numpy_input
)

Makes a tf.data.Dataset 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.

Example:

strategy = tf.distribute.MirroredStrategy()
numpy_input = np.ones([10], dtype=np.float32)
dataset = strategy.experimental_make_numpy_dataset(numpy_input)
dataset
<TensorSliceDataset shapes: (), types: tf.float32>
dataset = dataset.batch(2)
dist_dataset = strategy.experimental_distribute_dataset(dataset)

experimental_replicate_to_logical_devices

View source

experimental_replicate_to_logical_devices(
    tensor
)

Adds annotation that tensor will be replicated to all logical devices.

Note: This API is only supported in TPUStrategy for now. This adds an annotation to tensor tensor specifying that operations on tensor will be invoked on all logical devices.

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

iterator = iter(inputs)

@tf.function()
def step_fn(inputs):
  images, labels = inputs
  images = strategy.experimental_split_to_logical_devices(
    inputs, [1, 2, 4, 1])

  # model() function will be executed on 8 logical devices with `inputs`
  # split 2 * 4  ways.
  output = model(inputs)

  # For loss calculation, all logical devices share the same logits
  # and labels.
  labels = strategy.experimental_replicate_to_logical_devices(labels)
  output = strategy.experimental_replicate_to_logical_devices(output)
  loss = loss_fn(labels, output)

  return loss

strategy.run(step_fn, args=(next(iterator),))

Args: tensor: Input tensor to annotate.

experimental_split_to_logical_devices

View source

experimental_split_to_logical_devices(
    tensor, partition_dimensions
)

Adds annotation that tensor will be split across logical devices.

Note: This API is only supported in TPUStrategy for now. This adds an annotation to tensor tensor specifying that operations on tensor will be be split among multiple logical devices. Tensor tensor will be split across dimensions specified by partition_dimensions. The dimensions of tensor must be divisible by corresponding value in partition_dimensions.

For example, for system with 8 logical devices, if tensor is an image tensor with shape (batch_size, width, height, channel) and partition_dimensions is [1, 2, 4, 1], then tensor will be split 2 in width dimension and 4 way in height dimension and the split tensor values will be fed into 8 logical devices.

# Initializing TPU system with 8 logical devices and 1 replica.
resolver = tf.distribute.cluster_resolver.TPUClusterResolver(tpu='')
tf.config.experimental_connect_to_cluster(resolver)
topology = tf.tpu.experimental.initialize_tpu_system(resolver)
device_assignment = tf.tpu.experimental.DeviceAssignment.build(
    topology,
    computation_shape=[1, 2, 2, 2],
    num_replicas=1)
strategy = tf.distribute.TPUStrategy(
    resolver, experimental_device_assignment=device_assignment)

iterator = iter(inputs)

@tf.function()
def step_fn(inputs):
  inputs = strategy.experimental_split_to_logical_devices(
    inputs, [1, 2, 4, 1])

  # model() function will be executed on 8 logical devices with `inputs`
  # split 2 * 4  ways.
  output = model(inputs)
  return output

strategy.run(step_fn, args=(next(iterator),))

Args: tensor: Input tensor to annotate. partition_dimensions: An unnested list of integers with the size equal to rank of tensor specifying how tensor will be partitioned. The product of all elements in partition_dimensions must be equal to the total number of logical devices per replica.

reduce

View source

reduce(
    reduce_op, value, axis
)

Reduce value across replicas.

In OneDeviceStrategy, there is only one replica, so if axis=None, value is simply returned. If axis is specified as something other than None, such as axis=0, value is reduced along that axis and returned.

Example:

t = tf.range(10)

result = strategy.reduce(tf.distribute.ReduceOp.SUM, t, axis=None).numpy()
# result: [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]

result = strategy.reduce(tf.distribute.ReduceOp.SUM, t, axis=0).numpy()
# result: 45

run

View source

run(
    fn, args=(), kwargs=None, options=None
)

Run fn on each replica, with the given arguments.

In OneDeviceStrategy, fn is simply called within a device scope for the given device, with the provided arguments.

scope

View source

scope()

Returns a context manager selecting this Strategy as current.

Inside a with strategy.scope(): code block, this thread will use a variable creator set by strategy, and will enter its "cross-replica context".

In OneDeviceStrategy, all variables created inside strategy.scope() will be on device specified at strategy construction time. See example in the docs for this class.