tf.distribute.OneDeviceStrategy

TensorFlow 1 version View source on GitHub

A distribution strategy for running on a single device.

Inherits From: Strategy

tf.distribute.OneDeviceStrategy(
    device
)

Used in the notebooks

Used in the tutorials

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

Args:

  • device: Device string identifier for the device on which the variables should be placed. See class docs for more details on how the device is used. Examples: "/cpu:0", "/gpu:0", "/device:CPU:0", "/device:GPU:0"

Attributes:

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

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, 2],
    num_replicas=4)
strategy = tf.distribute.experimental.TPUStrategy(
    resolver, 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),))

Args:

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

Raises:

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

Returns:

Annotated tensor with idential value as tensor.

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.

Returns:

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

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,))

IMPORTANT: The tf.data.Dataset returned by dataset_fn should have a per-replica batch size, unlike experimental_distribute_dataset, which uses the global batch size. This may be computed using input_context.get_per_replica_batch_size.

Args:

Returns:

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

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.

Args:

  • 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.

Returns:

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 = ( 
     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]

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.

Args:

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

Returns:

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

experimental_make_numpy_dataset

View source

experimental_make_numpy_dataset(
    numpy_input
)

Makes a tf.data.Dataset for input provided via a numpy array.

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:

numpy_input = np.ones([10], dtype=np.float32)
dataset = strategy.experimental_make_numpy_dataset(numpy_input)
dist_dataset = strategy.experimental_distribute_dataset(dataset)

Args:

  • numpy_input: A nest of NumPy input arrays that will be converted into a dataset. Note that lists of Numpy arrays are stacked, as that is normal tf.data.Dataset behavior.

Returns:

A tf.data.Dataset representing numpy_input.

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, 2],
    num_replicas=4)
strategy = tf.distribute.experimental.TPUStrategy(
    resolver, 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.

Returns:

Annotated tensor with idential value as tensor.

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=[2, 2, 2],
    num_replicas=1)
strategy = tf.distribute.experimental.TPUStrategy(
    resolver, 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.

Raises:

  • ValueError: 1) If the size of partition_dimensions does not equal to rank of tensor or 2) if product of elements of partition_dimensions does not match the number of logical devices per replica defined by the implementing DistributionStrategy's device specification or 3) if a known size of tensor is not divisible by corresponding value in partition_dimensions.

Returns:

Annotated tensor with idential value as tensor.

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

Args:

  • reduce_op: A tf.distribute.ReduceOp value specifying how values should be combined.
  • value: A "per replica" value, e.g. returned by run to be combined into a single tensor.
  • axis: Specifies the dimension to reduce along within each replica's tensor. Should typically be set to the batch dimension, or None to only reduce across replicas (e.g. if the tensor has no batch dimension).

Returns:

A Tensor.

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.

Args:

  • fn: The function to run. The output must be a tf.nest of Tensors.
  • args: (Optional) Positional arguments to fn.
  • kwargs: (Optional) Keyword arguments to fn.
  • options: (Optional) An instance of tf.distribute.RunOptions specifying the options to run fn.

Returns:

Return value from running fn.

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.

Returns:

A context manager to use for creating variables with this strategy.