A Dataset comprising lines from one or more text files.
Inherits From: Dataset, Dataset
View source
Applies a transformation function to this dataset.
apply enables chaining of custom Dataset transformations, which are represented as functions that take one Dataset argument and return a transformed Dataset.
View source
Returns an iterator which converts all elements of the dataset to numpy.
Use as_numpy_iterator to inspect the content of your dataset. To see element shapes and types, print dataset elements directly instead of using as_numpy_iterator.
This method requires that you are running in eager mode and the dataset's element_spec contains only TensorSpec components.
as_numpy_iterator() will preserve the nested structure of dataset elements.
View source
Combines consecutive elements of this dataset into batches.
The components of the resulting element will have an additional outer dimension, which will be batch_size (or N % batch_size for the last element if batch_size does not divide the number of input elements N evenly and drop_remainder is False). If your program depends on the batches having the same outer dimension, you should set the drop_remainder argument to True to prevent the smaller batch from being produced.
View source
Caches the elements in this dataset.
The first time the dataset is iterated over, its elements will be cached either in the specified file or in memory. Subsequent iterations will use the cached data.
When caching to a file, the cached data will persist across runs. Even the first iteration through the data will read from the cache file. Changing the input pipeline before the call to .cache() will have no effect until the cache file is removed or the filename is changed.
View source
Returns the cardinality of the dataset, if known.
cardinality may return tf.data.INFINITE_CARDINALITY if the dataset contains an infinite number of elements or tf.data.UNKNOWN_CARDINALITY if the analysis fails to determine the number of elements in the dataset (e.g. when the dataset source is a file).
View source
Creates a Dataset by concatenating the given dataset with this dataset.
View source
Enumerates the elements of this dataset.
It is similar to python's enumerate.
View source
Filters this dataset according to predicate.
View source
Filters this dataset according to predicate. (deprecated)
View source
Maps map_func across this dataset and flattens the result.
Use flat_map if you want to make sure that the order of your dataset stays the same. For example, to flatten a dataset of batches into a dataset of their elements:
tf.data.Dataset.interleave() is a generalization of flat_map, since flat_map produces the same output as tf.data.Dataset.interleave(cycle_length=1)
View source
Creates a Dataset whose elements are generated by generator. (deprecated arguments)
The generator argument must be a callable object that returns an object that supports the iter() protocol (e.g. a generator function).
The elements generated by generator must be compatible with either the given output_signature argument or with the given output_types and (optionally) output_shapes arguments, whichiver was specified.
The recommended way to call from_generator is to use the output_signature argument. In this case the output will be assumed to consist of objects with the classes, shapes and types defined by tf.TypeSpec objects from output_signature argument:
There is also a deprecated way to call from_generator by either with output_types argument alone or together with output_shapes argument. In this case the output of the function will be assumed to consist of tf.Tensor objects with with the types defined by output_types and with the shapes which are either unknown or defined by output_shapes.
View source
Splits each rank-N tf.sparse.SparseTensor in this dataset row-wise. (deprecated)
View source
Creates a Dataset whose elements are slices of the given tensors.
The given tensors are sliced along their first dimension. This operation preserves the structure of the input tensors, removing the first dimension of each tensor and using it as the dataset dimension. All input tensors must have the same size in their first dimensions.
Note that if tensors contains a NumPy array, and eager execution is not enabled, the values will be embedded in the graph as one or more tf.constant operations. For large datasets (> 1 GB), this can waste memory and run into byte limits of graph serialization. If tensors contains one or more large NumPy arrays, consider the alternative described in this guide.
View source
Creates a Dataset with a single element, comprising the given tensors.
from_tensors produces a dataset containing only a single element. To slice the input tensor into multiple elements, use from_tensor_slices instead.
Note that if tensors contains a NumPy array, and eager execution is not enabled, the values will be embedded in the graph as one or more tf.constant operations. For large datasets (> 1 GB), this can waste memory and run into byte limits of graph serialization. If tensors contains one or more large NumPy arrays, consider the alternative described in this guide.
View source
Maps map_func across this dataset, and interleaves the results.
For example, you can use Dataset.interleave() to process many input files concurrently:
The cycle_length and block_length arguments control the order in which elements are produced. cycle_length controls the number of input elements that are processed concurrently. If you set cycle_length to 1, this transformation will handle one input element at a time, and will produce identical results to tf.data.Dataset.flat_map. In general, this transformation will apply map_func to cycle_length input elements, open iterators on the returned Dataset objects, and cycle through them producing block_length consecutive elements from each iterator, and consuming the next input element each time it reaches the end of an iterator.
Performance can often be improved by setting num_parallel_calls so that interleave will use multiple threads to fetch elements. If determinism isn't required, it can also improve performance to set deterministic=False.
View source
A dataset of all files matching one or more glob patterns.
The file_pattern argument should be a small number of glob patterns. If your filenames have already been globbed, use Dataset.from_tensor_slices(filenames) instead, as re-globbing every filename with list_files may result in poor performance with remote storage systems.
If we had the following files on our filesystem:
If we pass "/path/to/dir/*.py" as the directory, the dataset would produce:
View source
Creates an iterator for elements of this dataset. (deprecated)
View source
Creates an iterator for elements of this dataset. (deprecated)
View source
Maps map_func across the elements of this dataset.
This transformation applies map_func to each element of this dataset, and returns a new dataset containing the transformed elements, in the same order as they appeared in the input. map_func can be used to change both the values and the structure of a dataset's elements. For example, adding 1 to each element, or projecting a subset of element components.
The input signature of map_func is determined by the structure of each element in this dataset.
The value or values returned by map_func determine the structure of each element in the returned dataset.
map_func can accept as arguments and return any type of dataset element.
Note that irrespective of the context in which map_func is defined (eager vs. graph), tf.data traces the function and executes it as a graph. To use Python code inside of the function you have a few options:
1) Rely on AutoGraph to convert Python code into an equivalent graph computation. The downside of this approach is that AutoGraph can convert some but not all Python code.
2) Use tf.py_function, which allows you to write arbitrary Python code but will generally result in worse performance than 1). For example:
3) Use tf.numpy_function, which also allows you to write arbitrary Python code. Note that tf.py_function accepts tf.Tensor whereas tf.numpy_function accepts numpy arrays and returns only numpy arrays. For example:
Note that the use of tf.numpy_function and tf.py_function in general precludes the possibility of executing user-defined transformations in parallel (because of Python GIL).
Performance can often be improved by setting num_parallel_calls so that map will use multiple threads to process elements. If deterministic order isn't required, it can also improve performance to set deterministic=False.
View source
Maps map_func across the elements of this dataset. (deprecated)
View source
Returns the options for this dataset and its inputs.
View source
Combines consecutive elements of this dataset into padded batches.
This transformation combines multiple consecutive elements of the input dataset into a single element.
Like tf.data.Dataset.batch, the components of the resulting element will have an additional outer dimension, which will be batch_size (or N % batch_size for the last element if batch_size does not divide the number of input elements N evenly and drop_remainder is False). If your program depends on the batches having the same outer dimension, you should set the drop_remainder argument to True to prevent the smaller batch from being produced.
Unlike tf.data.Dataset.batch, the input elements to be batched may have different shapes, and this transformation will pad each component to the respective shape in padded_shapes. The padded_shapes argument determines the resulting shape for each dimension of each component in an output element:
See also tf.data.experimental.dense_to_sparse_batch, which combines elements that may have different shapes into a tf.sparse.SparseTensor.
View source
Creates a Dataset that prefetches elements from this dataset.
Most dataset input pipelines should end with a call to prefetch. This allows later elements to be prepared while the current element is being processed. This often improves latency and throughput, at the cost of using additional memory to store prefetched elements.
View source
Creates a Dataset of a step-separated range of values.
View source
Reduces the input dataset to a single element.
The transformation calls reduce_func successively on every element of the input dataset until the dataset is exhausted, aggregating information in its internal state. The initial_state argument is used for the initial state and the final state is returned as the result.
View source
Repeats this dataset so each original value is seen count times.
View source
Creates a Dataset that includes only 1/num_shards of this dataset.
shard is deterministic. The Dataset produced by A.shard(n, i) will contain all elements of A whose index mod n = i.
This dataset operator is very useful when running distributed training, as it allows each worker to read a unique subset.
When reading a single input file, you can shard elements as follows:
View source
Randomly shuffles the elements of this dataset.
This dataset fills a buffer with buffer_size elements, then randomly samples elements from this buffer, replacing the selected elements with new elements. For perfect shuffling, a buffer size greater than or equal to the full size of the dataset is required.
For instance, if your dataset contains 10,000 elements but buffer_size is set to 1,000, then shuffle will initially select a random element from only the first 1,000 elements in the buffer. Once an element is selected, its space in the buffer is replaced by the next (i.e. 1,001-st) element, maintaining the 1,000 element buffer.
reshuffle_each_iteration controls whether the shuffle order should be different for each epoch. In TF 1.X, the idiomatic way to create epochs was through the repeat transformation:
In TF 2.0, tf.data.Dataset objects are Python iterables which makes it possible to also create epochs through Python iteration:
View source
Creates a Dataset that skips count elements from this dataset.
View source
Creates a Dataset with at most count elements from this dataset.
View source
Splits elements of a dataset into multiple elements.
For example, if elements of the dataset are shaped [B, a0, a1, ...], where B may vary for each input element, then for each element in the dataset, the unbatched dataset will contain B consecutive elements of shape [a0, a1, ...].
View source
Combines (nests of) input elements into a dataset of (nests of) windows.
A "window" is a finite dataset of flat elements of size size (or possibly fewer if there are not enough input elements to fill the window and drop_remainder evaluates to False).
The shift argument determines the number of input elements by which the window moves on each iteration. If windows and elements are both numbered starting at 0, the first element in window k will be element k * shift of the input dataset. In particular, the first element of the first window will always be the first element of the input dataset.
The stride argument determines the stride of the input elements, and the shift argument determines the shift of the window.
Note that when the window transformation is applied to a dataset of nested elements, it produces a dataset of nested windows.
View source
Returns a new tf.data.Dataset with the given options set.
The options are "global" in the sense they apply to the entire dataset. If options are set multiple times, they are merged as long as different options do not use different non-default values.
View source
Creates a Dataset by zipping together the given datasets.
This method has similar semantics to the built-in zip() function in Python, with the main difference being that the datasets argument can be an arbitrary nested structure of Dataset objects.
View source
View source
Creates an iterator for elements of this dataset.
The returned iterator implements the Python Iterator protocol.
View source
Returns the length of the dataset if it is known and finite.
This method requires that you are running in eager mode, and that the length of the dataset is known and non-infinite. When the length may be unknown or infinite, or if you are running in graph mode, use tf.data.Dataset.cardinality instead.
View source


A Dataset comprising records from one or more TFRecord files.
Inherits From: Dataset, Dataset
View source
Applies a transformation function to this dataset.
apply enables chaining of custom Dataset transformations, which are represented as functions that take one Dataset argument and return a transformed Dataset.
View source
Returns an iterator which converts all elements of the dataset to numpy.
Use as_numpy_iterator to inspect the content of your dataset. To see element shapes and types, print dataset elements directly instead of using as_numpy_iterator.
This method requires that you are running in eager mode and the dataset's element_spec contains only TensorSpec components.
as_numpy_iterator() will preserve the nested structure of dataset elements.
View source
Combines consecutive elements of this dataset into batches.
The components of the resulting element will have an additional outer dimension, which will be batch_size (or N % batch_size for the last element if batch_size does not divide the number of input elements N evenly and drop_remainder is False). If your program depends on the batches having the same outer dimension, you should set the drop_remainder argument to True to prevent the smaller batch from being produced.
View source
Caches the elements in this dataset.
The first time the dataset is iterated over, its elements will be cached either in the specified file or in memory. Subsequent iterations will use the cached data.
When caching to a file, the cached data will persist across runs. Even the first iteration through the data will read from the cache file. Changing the input pipeline before the call to .cache() will have no effect until the cache file is removed or the filename is changed.
View source
Returns the cardinality of the dataset, if known.
cardinality may return tf.data.INFINITE_CARDINALITY if the dataset contains an infinite number of elements or tf.data.UNKNOWN_CARDINALITY if the analysis fails to determine the number of elements in the dataset (e.g. when the dataset source is a file).
View source
Creates a Dataset by concatenating the given dataset with this dataset.
View source
Enumerates the elements of this dataset.
It is similar to python's enumerate.
View source
Filters this dataset according to predicate.
View source
Filters this dataset according to predicate. (deprecated)
View source
Maps map_func across this dataset and flattens the result.
Use flat_map if you want to make sure that the order of your dataset stays the same. For example, to flatten a dataset of batches into a dataset of their elements:
tf.data.Dataset.interleave() is a generalization of flat_map, since flat_map produces the same output as tf.data.Dataset.interleave(cycle_length=1)
View source
Creates a Dataset whose elements are generated by generator. (deprecated arguments)
The generator argument must be a callable object that returns an object that supports the iter() protocol (e.g. a generator function).
The elements generated by generator must be compatible with either the given output_signature argument or with the given output_types and (optionally) output_shapes arguments, whichiver was specified.
The recommended way to call from_generator is to use the output_signature argument. In this case the output will be assumed to consist of objects with the classes, shapes and types defined by tf.TypeSpec objects from output_signature argument:
There is also a deprecated way to call from_generator by either with output_types argument alone or together with output_shapes argument. In this case the output of the function will be assumed to consist of tf.Tensor objects with with the types defined by output_types and with the shapes which are either unknown or defined by output_shapes.
View source
Splits each rank-N tf.sparse.SparseTensor in this dataset row-wise. (deprecated)
View source
Creates a Dataset whose elements are slices of the given tensors.
The given tensors are sliced along their first dimension. This operation preserves the structure of the input tensors, removing the first dimension of each tensor and using it as the dataset dimension. All input tensors must have the same size in their first dimensions.
Note that if tensors contains a NumPy array, and eager execution is not enabled, the values will be embedded in the graph as one or more tf.constant operations. For large datasets (> 1 GB), this can waste memory and run into byte limits of graph serialization. If tensors contains one or more large NumPy arrays, consider the alternative described in this guide.
View source
Creates a Dataset with a single element, comprising the given tensors.
from_tensors produces a dataset containing only a single element. To slice the input tensor into multiple elements, use from_tensor_slices instead.
Note that if tensors contains a NumPy array, and eager execution is not enabled, the values will be embedded in the graph as one or more tf.constant operations. For large datasets (> 1 GB), this can waste memory and run into byte limits of graph serialization. If tensors contains one or more large NumPy arrays, consider the alternative described in this guide.
View source
Maps map_func across this dataset, and interleaves the results.
For example, you can use Dataset.interleave() to process many input files concurrently:
The cycle_length and block_length arguments control the order in which elements are produced. cycle_length controls the number of input elements that are processed concurrently. If you set cycle_length to 1, this transformation will handle one input element at a time, and will produce identical results to tf.data.Dataset.flat_map. In general, this transformation will apply map_func to cycle_length input elements, open iterators on the returned Dataset objects, and cycle through them producing block_length consecutive elements from each iterator, and consuming the next input element each time it reaches the end of an iterator.
Performance can often be improved by setting num_parallel_calls so that interleave will use multiple threads to fetch elements. If determinism isn't required, it can also improve performance to set deterministic=False.
View source
A dataset of all files matching one or more glob patterns.
The file_pattern argument should be a small number of glob patterns. If your filenames have already been globbed, use Dataset.from_tensor_slices(filenames) instead, as re-globbing every filename with list_files may result in poor performance with remote storage systems.
If we had the following files on our filesystem:
If we pass "/path/to/dir/*.py" as the directory, the dataset would produce:
View source
Creates an iterator for elements of this dataset. (deprecated)
View source
Creates an iterator for elements of this dataset. (deprecated)
View source
Maps map_func across the elements of this dataset.
This transformation applies map_func to each element of this dataset, and returns a new dataset containing the transformed elements, in the same order as they appeared in the input. map_func can be used to change both the values and the structure of a dataset's elements. For example, adding 1 to each element, or projecting a subset of element components.
The input signature of map_func is determined by the structure of each element in this dataset.
The value or values returned by map_func determine the structure of each element in the returned dataset.
map_func can accept as arguments and return any type of dataset element.
Note that irrespective of the context in which map_func is defined (eager vs. graph), tf.data traces the function and executes it as a graph. To use Python code inside of the function you have a few options:
1) Rely on AutoGraph to convert Python code into an equivalent graph computation. The downside of this approach is that AutoGraph can convert some but not all Python code.
2) Use tf.py_function, which allows you to write arbitrary Python code but will generally result in worse performance than 1). For example:
3) Use tf.numpy_function, which also allows you to write arbitrary Python code. Note that tf.py_function accepts tf.Tensor whereas tf.numpy_function accepts numpy arrays and returns only numpy arrays. For example:
Note that the use of tf.numpy_function and tf.py_function in general precludes the possibility of executing user-defined transformations in parallel (because of Python GIL).
Performance can often be improved by setting num_parallel_calls so that map will use multiple threads to process elements. If deterministic order isn't required, it can also improve performance to set deterministic=False.
View source
Maps map_func across the elements of this dataset. (deprecated)
View source
Returns the options for this dataset and its inputs.
View source
Combines consecutive elements of this dataset into padded batches.
This transformation combines multiple consecutive elements of the input dataset into a single element.
Like tf.data.Dataset.batch, the components of the resulting element will have an additional outer dimension, which will be batch_size (or N % batch_size for the last element if batch_size does not divide the number of input elements N evenly and drop_remainder is False). If your program depends on the batches having the same outer dimension, you should set the drop_remainder argument to True to prevent the smaller batch from being produced.
Unlike tf.data.Dataset.batch, the input elements to be batched may have different shapes, and this transformation will pad each component to the respective shape in padded_shapes. The padded_shapes argument determines the resulting shape for each dimension of each component in an output element:
See also tf.data.experimental.dense_to_sparse_batch, which combines elements that may have different shapes into a tf.sparse.SparseTensor.
View source
Creates a Dataset that prefetches elements from this dataset.
Most dataset input pipelines should end with a call to prefetch. This allows later elements to be prepared while the current element is being processed. This often improves latency and throughput, at the cost of using additional memory to store prefetched elements.
View source
Creates a Dataset of a step-separated range of values.
View source
Reduces the input dataset to a single element.
The transformation calls reduce_func successively on every element of the input dataset until the dataset is exhausted, aggregating information in its internal state. The initial_state argument is used for the initial state and the final state is returned as the result.
View source
Repeats this dataset so each original value is seen count times.
View source
Creates a Dataset that includes only 1/num_shards of this dataset.
shard is deterministic. The Dataset produced by A.shard(n, i) will contain all elements of A whose index mod n = i.
This dataset operator is very useful when running distributed training, as it allows each worker to read a unique subset.
When reading a single input file, you can shard elements as follows:
View source
Randomly shuffles the elements of this dataset.
This dataset fills a buffer with buffer_size elements, then randomly samples elements from this buffer, replacing the selected elements with new elements. For perfect shuffling, a buffer size greater than or equal to the full size of the dataset is required.
For instance, if your dataset contains 10,000 elements but buffer_size is set to 1,000, then shuffle will initially select a random element from only the first 1,000 elements in the buffer. Once an element is selected, its space in the buffer is replaced by the next (i.e. 1,001-st) element, maintaining the 1,000 element buffer.
reshuffle_each_iteration controls whether the shuffle order should be different for each epoch. In TF 1.X, the idiomatic way to create epochs was through the repeat transformation:
In TF 2.0, tf.data.Dataset objects are Python iterables which makes it possible to also create epochs through Python iteration:
View source
Creates a Dataset that skips count elements from this dataset.
View source
Creates a Dataset with at most count elements from this dataset.
View source
Splits elements of a dataset into multiple elements.
For example, if elements of the dataset are shaped [B, a0, a1, ...], where B may vary for each input element, then for each element in the dataset, the unbatched dataset will contain B consecutive elements of shape [a0, a1, ...].
View source
Combines (nests of) input elements into a dataset of (nests of) windows.
A "window" is a finite dataset of flat elements of size size (or possibly fewer if there are not enough input elements to fill the window and drop_remainder evaluates to False).
The shift argument determines the number of input elements by which the window moves on each iteration. If windows and elements are both numbered starting at 0, the first element in window k will be element k * shift of the input dataset. In particular, the first element of the first window will always be the first element of the input dataset.
The stride argument determines the stride of the input elements, and the shift argument determines the shift of the window.
Note that when the window transformation is applied to a dataset of nested elements, it produces a dataset of nested windows.
View source
Returns a new tf.data.Dataset with the given options set.
The options are "global" in the sense they apply to the entire dataset. If options are set multiple times, they are merged as long as different options do not use different non-default values.
View source
Creates a Dataset by zipping together the given datasets.
This method has similar semantics to the built-in zip() function in Python, with the main difference being that the datasets argument can be an arbitrary nested structure of Dataset objects.
View source
View source
Creates an iterator for elements of this dataset.
The returned iterator implements the Python Iterator protocol.
View source
Returns the length of the dataset if it is known and finite.
This method requires that you are running in eager mode, and that the length of the dataset is known and non-infinite. When the length may be unknown or infinite, or if you are running in graph mode, use tf.data.Dataset.cardinality instead.
View source


Public API for tf.debugging namespace.
experimental module: Public API for tf.debugging.experimental namespace.
Assert(...): Asserts that the given condition is true.
assert_all_finite(...): Assert that the tensor does not contain any NaN's or Inf's.
assert_equal(...): Assert the condition x == y holds element-wise.
assert_greater(...): Assert the condition x > y holds element-wise.
assert_greater_equal(...): Assert the condition x >= y holds element-wise.
assert_integer(...): Assert that x is of integer dtype.
assert_less(...): Assert the condition x < y holds element-wise.
assert_less_equal(...): Assert the condition x <= y holds element-wise.
assert_near(...): Assert the condition x and y are close element-wise.
assert_negative(...): Assert the condition x < 0 holds element-wise.
assert_non_negative(...): Assert the condition x >= 0 holds element-wise.
assert_non_positive(...): Assert the condition x <= 0 holds element-wise.
assert_none_equal(...): Assert the condition x != y holds element-wise.
assert_positive(...): Assert the condition x > 0 holds element-wise.
assert_proper_iterable(...): Static assert that values is a "proper" iterable.
assert_rank(...): Assert x has rank equal to rank.
assert_rank_at_least(...): Assert x has rank equal to rank or higher.
assert_rank_in(...): Assert x has rank in ranks.
assert_same_float_dtype(...): Validate and return float type based on tensors and dtype.
assert_scalar(...): Asserts that the given tensor is a scalar (i.e. zero-dimensional).
assert_shapes(...): Assert tensor shapes and dimension size relationships between tensors.
assert_type(...): Statically asserts that the given Tensor is of the specified type.
check_numerics(...): Checks a tensor for NaN and Inf values.
disable_check_numerics(...): Disable the eager/graph unified numerics checking mechanism.
enable_check_numerics(...): Enable tensor numerics checking in an eager/graph unified fashion.
get_log_device_placement(...): Get if device placements are logged.
is_finite(...): Returns which elements of x are finite.
is_inf(...): Returns which elements of x are Inf.
is_nan(...): Returns which elements of x are NaN.
is_non_decreasing(...): Returns True if x is non-decreasing.
is_numeric_tensor(...): Returns True if the elements of tensor are numbers.
is_strictly_increasing(...): Returns True if x is strictly increasing.
set_log_device_placement(...): Set if device placements should be logged.


Assert tensor shapes and dimension size relationships between tensors.
This Op checks that a collection of tensors shape relationships satisfies given constraints.
Example of adding a dependency to an operation:
If x, y, param or scalar does not have a shape that satisfies all specified constraints, message, as well as the first summarize entries of the first encountered violating tensor are printed, and InvalidArgumentError is raised.
Size entries in the specified shapes are checked against other entries by their hash, except:
If the first entry of a shape is ... (type Ellipsis) or '*' that indicates a variable number of outer dimensions of unspecified size, i.e. the constraint applies to the inner-most dimensions only.
Scalar tensors and specified shapes of length zero (excluding the 'inner-most' prefix) are both treated as having a single dimension of size one.


Public API for tf.debugging.experimental namespace.
disable_dump_debug_info(...): Disable the currently-enabled debugging dumping.
enable_dump_debug_info(...): Enable dumping debugging information from a TensorFlow program.


Convert CSV records to tensors. Each column maps to one tensor.
Compat aliases for migration
See Migration guide for more details.
tf.compat.v1.io.decode_csv
RFC 4180 format is expected for the CSV records. (https://tools.ietf.org/html/rfc4180) Note that we allow leading and trailing spaces with int or float field.


Convert raw byte strings into tensors. (deprecated arguments)
Compat aliases for migration
See Migration guide for more details.
tf.compat.v1.io.decode_raw


Delete the tensor for the given tensor handle.
This is EXPERIMENTAL and subject to change.
Delete the tensor of a given tensor handle. The tensor is produced in a previous run() and stored in the state of the session.


DepthToSpace for tensors of type T.
Compat aliases for migration
See Migration guide for more details.
tf.compat.v1.nn.depth_to_space
Rearranges data from depth into blocks of spatial data. This is the reverse transformation of SpaceToDepth. More specifically, this op outputs a copy of the input tensor where values from the depth dimension are moved in spatial blocks to the height and width dimensions. The attr block_size indicates the input block size and how the data is moved.
The data_format attr specifies the layout of the input and output tensors with the following options: "NHWC": [ batch, height, width, channels ] "NCHW": [ batch, channels, height, width ] "NCHW_VECT_C": qint8 [ batch, channels / 4, height, width, 4 ]
It is useful to consider the operation as transforming a 6-D Tensor. e.g. for data_format = NHWC, Each element in the input tensor can be specified via 6 coordinates, ordered by decreasing memory layout significance as: n,iY,iX,bY,bX,oC (where n=batch index, iX, iY means X or Y coordinates within the input image, bX, bY means coordinates within the output block, oC means output channels). The output would be the input transposed to the following layout: n,iY,bY,iX,bX,oC
This operation is useful for resizing the activations between convolutions (but keeping all data), e.g. instead of pooling. It is also useful for training purely convolutional models.
For example, given an input of shape [1, 1, 1, 4], data_format = "NHWC" and block_size = 2:
This operation will output a tensor of shape [1, 2, 2, 1]:
Here, the input has a batch of 1 and each batch element has shape [1, 1, 4], the corresponding output will have 2x2 elements and will have a depth of 1 channel (1 = 4 / (block_size * block_size)). The output element shape is [2, 2, 1].
For an input tensor with larger depth, here of shape [1, 1, 1, 12], e.g.
This operation, for block size of 2, will return the following tensor of shape [1, 2, 2, 3]
Similarly, for the following input of shape [1 2 2 4], and a block size of 2:
the operator will return the following tensor of shape [1 4 4 1]:


Wrapper for Graph.device() using the default graph.
See tf.Graph.device for more details.


Represents a (possibly partial) specification for a TensorFlow device.
Inherits From: DeviceSpec
DeviceSpecs are used throughout TensorFlow to describe where state is stored and computations occur. Using DeviceSpec allows you to parse device spec strings to verify their validity, merge them or compose them programmatically.
With eager execution disabled (by default in TensorFlow 1.x and by calling disable_eager_execution() in TensorFlow 2.x), the following syntax can be used:
A DeviceSpec consists of 5 components -- each of which is optionally specified:
View source
Construct a DeviceSpec from a string.
View source
Returns a new DeviceSpec which incorporates dev.
When combining specs, dev will take precedence over the current spec. So for instance:
is equivalent to:
View source
Merge the properties of "dev" into this DeviceSpec.
View source
Parse a DeviceSpec name into its components.
2.x behavior change: In TensorFlow 1.x, this function mutates its own state and returns itself. In 2.x, DeviceSpecs are immutable, and this function will return a DeviceSpec which contains the spec.
Recommended:
Will work in 1.x and 2.x (though deprecated in 2.x):
Will NOT work in 2.x:
In general, DeviceSpec.from_string should completely replace DeviceSpec.parse_from_string, and DeviceSpec.replace should completely replace setting attributes directly.
View source
Convenience method for making a new DeviceSpec by overriding fields.
View source
Return a string representation of this DeviceSpec.
View source
Checks if the other DeviceSpec is same as the current instance, eg have
same value for all the internal fields.


Represents the value of one dimension in a TensorShape.
View source
Raises an exception if other is not compatible with this Dimension.
View source
Returns true if other is compatible with this Dimension.
Two known Dimensions are compatible if they have the same value. An unknown Dimension is compatible with all other Dimensions.
View source
Returns a Dimension that combines the information in self and other.
Dimensions are combined as follows:
View source
Returns the sum of self and other.
Dimensions are summed as follows:
View source
DEPRECATED: Use __floordiv__ via x // y instead.
This function exists only for backwards compatibility purposes; new code should use __floordiv__ via the syntax x // y. Using x // y communicates clearly that the result rounds down, and is forward compatible to Python 3.
View source
Returns true if other has the same known value as this Dimension.
View source
Returns the quotient of self and other rounded down.
Dimensions are divided as follows:
View source
Returns True if self is known to be greater than or equal to other.
Dimensions are compared as follows:
View source
Returns True if self is known to be greater than other.
Dimensions are compared as follows:
View source
Returns True if self is known to be less than or equal to other.
Dimensions are compared as follows:
View source
Returns True if self is known to be less than other.
Dimensions are compared as follows:
View source
Returns self modulo other.
Dimension modulo are computed as follows:
View source
Returns the product of self and other.
Dimensions are summed as follows:
View source
Returns true if other has a different known value from self.
View source
Returns the sum of other and self.
View source
Use __floordiv__ via x // y instead.
This function exists only to have a better error message. Instead of: TypeError: unsupported operand type(s) for /: 'int' and 'Dimension', this function will explicitly call for usage of // instead.
View source
Returns the quotient of other and self rounded down.
View source
Returns other modulo self.
View source
Returns the product of self and other.
View source
Returns the subtraction of self from other.
View source
Use __floordiv__ via x // y instead.
This function exists only to have a better error message. Instead of: TypeError: unsupported operand type(s) for /: 'int' and 'Dimension', this function will explicitly call for usage of // instead.
View source
Returns the subtraction of other from self.
Dimensions are subtracted as follows:
View source
Use __floordiv__ via x // y instead.
This function exists only to have a better error message. Instead of: TypeError: unsupported operand type(s) for /: 'Dimension' and 'int', this function will explicitly call for usage of // instead.


Opts out of control flow v2.
If your code needs tf.disable_control_flow_v2() to be called to work properly please file a bug.


Disables eager execution.
This function can only be called before any Graphs, Ops, or Tensors have been created. It can be used at the beginning of the program for complex migration projects from TensorFlow 1.x to 2.x.


Opts out of resource variables. (deprecated)
If your code needs tf.disable_resource_variables() to be called to work properly please file a bug.


Compare Tensors by their id and be hashable.
This is a legacy behaviour of TensorFlow and is highly discouraged.


Disables TensorFlow 2.x behaviors.
This function can be called at the beginning of the program (before Tensors, Graphs or other structures have been created, and before devices have been initialized. It switches all global behaviors that are different between TensorFlow 1.x and 2.x to behave as intended for 1.x.
User can call this function to disable 2.x behavior during complex migrations.


Disables the V2 TensorShape behavior and reverts to V1 behavior.
See docstring for enable_v2_tensorshape for details about the new behavior.


Library for running a computation across multiple devices.
The intent of this library is that you can write an algorithm in a stylized way and it will be usable with a variety of different tf.distribute.Strategy implementations. Each descendant will implement a different strategy for distributing the algorithm across multiple devices/machines. Furthermore, these changes can be hidden inside the specific layers and other library classes that need special treatment to run in a distributed setting, so that most users' model definition code can run unchanged. The tf.distribute.Strategy API works the same way with eager and graph execution.
Guides
Tutorials
Distributed Training Tutorials
The tutorials cover how to use tf.distribute.Strategy to do distributed training with native Keras APIs, custom training loops, and Esitmator APIs. They also cover how to save/load model when using tf.distribute.Strategy.
Glossary
Parameter servers: These are machines that hold a single copy of parameters/variables, used by some strategies (right now just tf.distribute.experimental.ParameterServerStrategy). All replicas that want to operate on a variable retrieve it at the beginning of a step and send an update to be applied at the end of the step. These can in priniciple support either sync or async training, but right now we only have support for async training with parameter servers. Compare to tf.distribute.experimental.CentralStorageStrategy, which puts all variables on a single device on the same machine (and does sync training), and tf.distribute.MirroredStrategy, which mirrors variables to multiple devices (see below).
Replica context vs. Cross-replica context vs Update context
A replica context applies when you execute the computation function that was called with strategy.run. Conceptually, you're in replica context when executing the computation function that is being replicated.
An update context is entered in a tf.distribute.StrategyExtended.update call.
An cross-replica context is entered when you enter a strategy.scope. This is useful for calling tf.distribute.Strategy methods which operate across the replicas (like reduce_to()). By default you start in a replica context (the "default single replica context") and then some methods can switch you back and forth.
Distributed value: Distributed value is represented by the base class tf.distribute.DistributedValues. tf.distribute.DistributedValues is useful to represent values on multiple devices, and it contains a map from replica id to values. Two representative kinds of tf.distribute.DistributedValues are "PerReplica" and "Mirrored" values.
"PerReplica" values exist on the worker devices, with a different value for each replica. They are produced by iterating through a distributed dataset returned by tf.distribute.Strategy.experimental_distribute_dataset and tf.distribute.Strategy.distribute_datasets_from_function. They are also the typical result returned by tf.distribute.Strategy.run.
"Mirrored" values are like "PerReplica" values, except we know that the value on all replicas are the same. We can safely read a "Mirrored" value in a cross-replica context by using the value on any replica.
Unwrapping and merging: Consider calling a function fn on multiple replicas, like strategy.run(fn, args=[w]) with an argument w that is a tf.distribute.DistributedValues. This means w will have a map taking replica id 0 to w0, replica id 1 to w1, etc. strategy.run() unwraps w before calling fn, so it calls fn(w0) on device d0, fn(w1) on device d1, etc. It then merges the return values from fn(), which leads to one common object if the returned values are the same object from every replica, or a DistributedValues object otherwise.
Reductions and all-reduce: A reduction is a method of aggregating multiple values into one value, like "sum" or "mean". If a strategy is doing sync training, we will perform a reduction on the gradients to a parameter from all replicas before applying the update. All-reduce is an algorithm for performing a reduction on values from multiple devices and making the result available on all of those devices.
Mirrored variables: These are variables that are created on multiple devices, where we keep the variables in sync by applying the same updates to every copy. Mirrored variables are created with tf.Variable(...synchronization=tf.VariableSynchronization.ON_WRITE...). Normally they are only used in synchronous training.
SyncOnRead variables
SyncOnRead variables are created by tf.Variable(...synchronization=tf.VariableSynchronization.ON_READ...), and they are created on multiple devices. In replica context, each component variable on the local replica can perform reads and writes without synchronization with each other. When the SyncOnRead variable is read in cross-replica context, the values from component variables are aggregated and returned.
SyncOnRead variables bring a lot of custom configuration difficulty to the underlying logic, so we do not encourage users to instantiate and use SyncOnRead variable on their own. We have mainly used SyncOnRead variables for use cases such as batch norm and metrics. For performance reasons, we often don't need to keep these statistics in sync every step and they can be accumulated on each replica independently. The only time we want to sync them is reporting or checkpointing, which typically happens in cross-replica context. SyncOnRead variables are also often used by advanced users who want to control when variable values are aggregated. For example, users sometimes want to maintain gradients independently on each replica for a couple of steps without aggregation.
Distribute-aware layers
Layers are generally called in a replica context, except when defining a Keras functional model. tf.distribute.in_cross_replica_context will let you determine which case you are in. If in a replica context, the tf.distribute.get_replica_context function will return the default replica context outside a strategy scope, None within a strategy scope, and a tf.distribute.ReplicaContext object inside a strategy scope and within a tf.distribute.Strategy.run function. The ReplicaContext object has an all_reduce method for aggregating across all replicas.
Note that we provide a default version of tf.distribute.Strategy that is used when no other strategy is in scope, that provides the same API with reasonable default behavior.
cluster_resolver module: Library imports for ClusterResolvers.
experimental module: Public API for tf.distribute.experimental namespace.
class CrossDeviceOps: Base class for cross-device reduction and broadcasting algorithms.
class HierarchicalCopyAllReduce: Hierarchical copy all-reduce implementation of CrossDeviceOps.
class InputContext: A class wrapping information needed by an input function.
class InputReplicationMode: Replication mode for input function.
class MirroredStrategy: Synchronous training across multiple replicas on one machine.
class NcclAllReduce: NCCL all-reduce implementation of CrossDeviceOps.
class OneDeviceStrategy: A distribution strategy for running on a single device.
class ReduceOp: Indicates how a set of values should be reduced.
class ReductionToOneDevice: A CrossDeviceOps implementation that copies values to one device to reduce.
class ReplicaContext: A class with a collection of APIs that can be called in a replica context.
class RunOptions: Run options for strategy.run.
class Server: An in-process TensorFlow server, for use in distributed training.
class Strategy: A list of devices with a state & compute distribution policy.
class StrategyExtended: Additional APIs for algorithms that need to be distribution-aware.
experimental_set_strategy(...): Set a tf.distribute.Strategy as current without with strategy.scope().
get_loss_reduction(...): tf.distribute.ReduceOp corresponding to the last loss reduction.
get_replica_context(...): Returns the current tf.distribute.ReplicaContext or None.
get_strategy(...): Returns the current tf.distribute.Strategy object.
has_strategy(...): Return if there is a current non-default tf.distribute.Strategy.
in_cross_replica_context(...): Returns True if in a cross-replica context.


Library imports for ClusterResolvers.
This library contains all implementations of ClusterResolvers. ClusterResolvers are a way of specifying cluster information for distributed execution. Built on top of existing ClusterSpec framework, ClusterResolvers are a way for TensorFlow to communicate with various cluster management systems (e.g. GCE, AWS, etc...).
class ClusterResolver: Abstract class for all implementations of ClusterResolvers.
class GCEClusterResolver: ClusterResolver for Google Compute Engine.
class KubernetesClusterResolver: ClusterResolver for Kubernetes.
class SimpleClusterResolver: Simple implementation of ClusterResolver that accepts all attributes.
class SlurmClusterResolver: ClusterResolver for system with Slurm workload manager.
class TFConfigClusterResolver: Implementation of a ClusterResolver which reads the TF_CONFIG EnvVar.
class TPUClusterResolver: Cluster Resolver for Google Cloud TPUs.
class UnionResolver: Performs a union on underlying ClusterResolvers.


Public API for tf.distribute.experimental namespace.
class CentralStorageStrategy: A one-machine strategy that puts all variables on a single device.
class CollectiveCommunication: Cross device communication implementation.
class CollectiveHints: Hints for collective operations like AllReduce.
class CommunicationImplementation: Cross device communication implementation.
class CommunicationOptions: Options for cross device communications like All-reduce.
class MultiWorkerMirroredStrategy: A distribution strategy for synchronous training on multiple workers.
class ParameterServerStrategy: An asynchronous multi-worker parameter server tf.distribute strategy.
class TPUStrategy: TPU distribution strategy implementation.


A one-machine strategy that puts all variables on a single device.
Inherits From: Strategy
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.
In general, when using a multi-worker tf.distribute strategy such as tf.distribute.experimental.MultiWorkerMirroredStrategy or tf.distribute.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,
For more information, please see tf.distribute.cluster_resolver.ClusterResolver's API docstring.
View source
Distributes tf.data.Dataset instances created by calls to dataset_fn.
The argument dataset_fn that users pass in is an input function that has a tf.distribute.InputContext argument and returns a tf.data.Dataset instance. It is expected that the returned dataset from dataset_fn is already batched by per-replica batch size (i.e. global batch size divided by the number of replicas in sync) and sharded. tf.distribute.Strategy.distribute_datasets_from_function does not batch or shard the tf.data.Dataset instance returned from the input function. dataset_fn will be called on the CPU device of each of the workers and each generates a dataset where every replica on that worker will dequeue one batch of inputs (i.e. if a worker has two replicas, two batches will be dequeued from the Dataset every step).
This method can be used for several purposes. First, it allows you to specify your own batching and sharding logic. (In contrast, tf.distribute.experimental_distribute_dataset does batching and sharding for you.) For example, where experimental_distribute_dataset is unable to shard the input files, this method might be used to manually shard the dataset (avoiding the slow fallback behavior in experimental_distribute_dataset). In cases where the dataset is infinite, this sharding can be done by creating dataset replicas that differ only in their random seed.
The dataset_fn should take an tf.distribute.InputContext instance where information about batching and input replication can be accessed.
You can use element_spec property of the tf.distribute.DistributedDataset returned by this API to query the tf.TypeSpec of the elements returned by the iterator. This can be used to set the input_signature property of a tf.function. Follow tf.distribute.DistributedDataset.element_spec to see an example.
For a tutorial on more usage and properties of this method, refer to the tutorial on distributed input). If you are interested in last partial batch handling, read this section.
View source
Creates tf.distribute.DistributedDataset from tf.data.Dataset.
The returned tf.distribute.DistributedDataset can be iterated over similar to regular datasets. NOTE: The user cannot add any more transformations to a tf.distribute.DistributedDataset. You can only create an iterator or examine the tf.TypeSpec of the data generated by it. See API docs of tf.distribute.DistributedDataset to learn more.
The following is an example:
Three key actions happending under the hood of this method are batching, sharding, and prefetching.
In the code snippet above, dataset is batched by global_batch_size, and calling experimental_distribute_dataset on it rebatches dataset to a new batch size that is equal to the global batch size divided by the number of replicas in sync. We iterate through it using a Pythonic for loop. x is a tf.distribute.DistributedValues containing data for all replicas, and each replica gets data of the new batch size. tf.distribute.Strategy.run will take care of feeding the right per-replica data in x to the right replica_fn executed on each replica.
Sharding contains autosharding across multiple workers and within every worker. First, in multi-worker distributed training (i.e. when you use tf.distribute.experimental.MultiWorkerMirroredStrategy or tf.distribute.TPUStrategy), autosharding a dataset over a set of workers means that each worker is assigned a subset of the entire dataset (if the right tf.data.experimental.AutoShardPolicy is set). This is to ensure that at each step, a global batch size of non-overlapping dataset elements will be processed by each worker. Autosharding has a couple of different options that can be specified using tf.data.experimental.DistributeOptions. Then, sharding within each worker means the method will split the data among all the worker devices (if more than one a present). This will happen regardless of multi-worker autosharding.
By default, this method adds a prefetch transformation at the end of the user provided tf.data.Dataset instance. The argument to the prefetch transformation which is buffer_size is equal to the number of replicas in sync.
If the above batch splitting and dataset sharding logic is undesirable, please use tf.distribute.Strategy.distribute_datasets_from_function instead, which does not do any automatic batching or sharding for you.
For a tutorial on more usage and properties of this method, refer to the tutorial on distributed input. If you are interested in last partial batch handling, read this section.
View source
Returns the list of all local per-replica values contained in value.
View source
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 tf.distribute.Strategy.experimental_distribute_dataset with the returned dataset to further distribute it with the strategy.
View source
Runs ops in fn on each replica, with inputs from input_iterator.
DEPRECATED: This method is not available in TF 2.x. Please switch to using run instead.
When eager execution is enabled, executes ops specified by fn on each replica. Otherwise, builds a graph to execute the ops on each replica.
Each replica will take a single, different input from the inputs provided by one get_next call on the input iterator.
fn may call tf.distribute.get_replica_context() to access members such as replica_id_in_sync_group.
View source
Makes an iterator for input provided via dataset.
DEPRECATED: This method is not available in TF 2.x.
Data from the given dataset will be distributed evenly across all the compute replicas. We will assume that the input dataset is batched by the global batch size. With this assumption, we will make a best effort to divide each batch across all the replicas (one or more workers). If this effort fails, an error will be thrown, and the user should instead use make_input_fn_iterator which provides more control to the user, and does not try to divide a batch across replicas.
The user could also use make_input_fn_iterator if they want to customize which input is fed to which replica/worker etc.
View source
Returns an iterator split across replicas created from an input function.
DEPRECATED: This method is not available in TF 2.x.
The input_fn should take an tf.distribute.InputContext object where information about batching and input sharding can be accessed:
The tf.data.Dataset returned by input_fn should have a per-replica batch size, which may be computed using input_context.get_per_replica_batch_size.
View source
Reduce value across replicas and return result on current device.
To see how this would look with multiple replicas, consider the same example with MirroredStrategy with 2 GPUs:
This API is typically used for aggregating the results returned from different replicas, for reporting etc. For example, loss computed from different replicas can be averaged using this API before printing.
There are a number of different tf.distribute APIs for reducing values across replicas:
What should axis be?
Given a per-replica value returned by run, say a per-example loss, the batch will be divided across all the replicas. This function allows you to aggregate across replicas and optionally also across batch elements by specifying the axis parameter accordingly.
For example, if you have a global batch size of 8 and 2 replicas, values for examples [0, 1, 2, 3] will be on replica 0 and [4, 5, 6, 7] will be on replica 1. With axis=None, reduce will aggregate only across replicas, returning [0+4, 1+5, 2+6, 3+7]. This is useful when each replica is computing a scalar or some other value that doesn't have a "batch" dimension (like a gradient or loss).
Sometimes, you will want to aggregate across both the global batch and all replicas. You can get this behavior by specifying the batch dimension as the axis, typically axis=0. In this case it would return a scalar 0+1+2+3+4+5+6+7.
If there is a last partial batch, you will need to specify an axis so that the resulting shape is consistent across replicas. So if the last batch has size 6 and it is divided into [0, 1, 2, 3] and [4, 5], you would get a shape mismatch unless you specify axis=0. If you specify tf.distribute.ReduceOp.MEAN, using axis=0 will use the correct denominator of 6. Contrast this with computing reduce_mean to get a scalar value on each replica and this function to average those means, which will weigh some values 1/8 and others 1/4.
View source
Invokes fn on each replica, with the given arguments.
This method is the primary way to distribute your computation with a tf.distribute object. It invokes fn on each replica. If args or kwargs have tf.distribute.DistributedValues, such as those produced by a tf.distribute.DistributedDataset from tf.distribute.Strategy.experimental_distribute_dataset or tf.distribute.Strategy.distribute_datasets_from_function, when fn is executed on a particular replica, it will be executed with the component of tf.distribute.DistributedValues that correspond to that replica.
fn is invoked under a replica context. fn may call tf.distribute.get_replica_context() to access members such as all_reduce. Please see the module-level docstring of tf.distribute for the concept of replica context.
All arguments in args or kwargs should either be Python values of a nested structure of tensors, e.g. a list of tensors, in which case args and kwargs will be passed to the fn invoked on each replica. Or args or kwargs can be tf.distribute.DistributedValues containing tensors or composite tensors, i.e. tf.compat.v1.TensorInfo.CompositeTensor, in which case each fn call will get the component of a tf.distribute.DistributedValues corresponding to its replica.
View source
Context manager to make the strategy current and distribute variables.
This method returns a context manager, and is used as follows:
What happens when Strategy.scope is entered?
What should be in scope and what should be outside?
There are a number of requirements on what needs to happen inside the scope. However, in places where we have information about which strategy is in use, we often enter the scope for the user, so they don't have to do it explicitly (i.e. calling those either inside or outside the scope is OK).
View source
Returns a copy of config_proto modified for use with this strategy.
DEPRECATED: This method is not available in TF 2.x.
The updated config has something needed to run a strategy, e.g. configuration to run collective ops, or device filters to improve distributed training performance.


A distribution strategy for synchronous training on multiple workers.
Inherits From: Strategy
This strategy implements synchronous distributed training across multiple workers, each with potentially multiple GPUs. Similar to tf.distribute.MirroredStrategy, it replicates all variables and computations to each local device. The difference is that it uses a distributed collective implementation (e.g. all-reduce), so that multiple workers can work together.
You need to launch your program on each worker and configure cluster_resolver correctly. For example, if you are using tf.distribute.cluster_resolver.TFConfigClusterResolver, each worker needs to have its corresponding task_type and task_id set in the TF_CONFIG environment variable. An example TF_CONFIG on worker-0 of a two worker cluster is:
Your program runs on each worker as-is. Note that collectives require each worker to participate. All tf.distribute and non tf.distribute API may use collectives internally, e.g. checkpointing and saving since reading a tf.Variable with tf.VariableSynchronization.ON_READ all-reduces the value. Therefore it's recommended to run exactly the same program on each worker. Dispatching based on task_type or task_id of the worker is error-prone.
cluster_resolver.num_accelerators() determines the number of GPUs the strategy uses. If it's zero, the strategy uses the CPU. All workers need to use the same number of devices, otherwise the behavior is undefined.
This strategy is not intended for TPU. Use tf.distribute.TPUStrategy instead.
After setting up TF_CONFIG, using this strategy is similar to using tf.distribute.MirroredStrategy and tf.distribute.TPUStrategy.
You can also write your own training loop:
See Multi-worker training with Keras for a detailed tutorial.
Saving
You need to save and checkpoint on all workers instead of just one. This is because variables whose synchronization=ON_READ triggers aggregation during saving. It's recommended to save to a different path on each worker to avoid race conditions. Each worker saves the same thing. See Multi-worker training with Keras tutorial for examples.
Known Issues
In general, when using a multi-worker tf.distribute strategy such as tf.distribute.experimental.MultiWorkerMirroredStrategy or tf.distribute.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,
For more information, please see tf.distribute.cluster_resolver.ClusterResolver's API docstring.
View source
Distributes tf.data.Dataset instances created by calls to dataset_fn.
The argument dataset_fn that users pass in is an input function that has a tf.distribute.InputContext argument and returns a tf.data.Dataset instance. It is expected that the returned dataset from dataset_fn is already batched by per-replica batch size (i.e. global batch size divided by the number of replicas in sync) and sharded. tf.distribute.Strategy.distribute_datasets_from_function does not batch or shard the tf.data.Dataset instance returned from the input function. dataset_fn will be called on the CPU device of each of the workers and each generates a dataset where every replica on that worker will dequeue one batch of inputs (i.e. if a worker has two replicas, two batches will be dequeued from the Dataset every step).
This method can be used for several purposes. First, it allows you to specify your own batching and sharding logic. (In contrast, tf.distribute.experimental_distribute_dataset does batching and sharding for you.) For example, where experimental_distribute_dataset is unable to shard the input files, this method might be used to manually shard the dataset (avoiding the slow fallback behavior in experimental_distribute_dataset). In cases where the dataset is infinite, this sharding can be done by creating dataset replicas that differ only in their random seed.
The dataset_fn should take an tf.distribute.InputContext instance where information about batching and input replication can be accessed.
You can use element_spec property of the tf.distribute.DistributedDataset returned by this API to query the tf.TypeSpec of the elements returned by the iterator. This can be used to set the input_signature property of a tf.function. Follow tf.distribute.DistributedDataset.element_spec to see an example.
For a tutorial on more usage and properties of this method, refer to the tutorial on distributed input). If you are interested in last partial batch handling, read this section.
View source
Creates tf.distribute.DistributedDataset from tf.data.Dataset.
The returned tf.distribute.DistributedDataset can be iterated over similar to regular datasets. NOTE: The user cannot add any more transformations to a tf.distribute.DistributedDataset. You can only create an iterator or examine the tf.TypeSpec of the data generated by it. See API docs of tf.distribute.DistributedDataset to learn more.
The following is an example:
Three key actions happending under the hood of this method are batching, sharding, and prefetching.
In the code snippet above, dataset is batched by global_batch_size, and calling experimental_distribute_dataset on it rebatches dataset to a new batch size that is equal to the global batch size divided by the number of replicas in sync. We iterate through it using a Pythonic for loop. x is a tf.distribute.DistributedValues containing data for all replicas, and each replica gets data of the new batch size. tf.distribute.Strategy.run will take care of feeding the right per-replica data in x to the right replica_fn executed on each replica.
Sharding contains autosharding across multiple workers and within every worker. First, in multi-worker distributed training (i.e. when you use tf.distribute.experimental.MultiWorkerMirroredStrategy or tf.distribute.TPUStrategy), autosharding a dataset over a set of workers means that each worker is assigned a subset of the entire dataset (if the right tf.data.experimental.AutoShardPolicy is set). This is to ensure that at each step, a global batch size of non-overlapping dataset elements will be processed by each worker. Autosharding has a couple of different options that can be specified using tf.data.experimental.DistributeOptions. Then, sharding within each worker means the method will split the data among all the worker devices (if more than one a present). This will happen regardless of multi-worker autosharding.
By default, this method adds a prefetch transformation at the end of the user provided tf.data.Dataset instance. The argument to the prefetch transformation which is buffer_size is equal to the number of replicas in sync.
If the above batch splitting and dataset sharding logic is undesirable, please use tf.distribute.Strategy.distribute_datasets_from_function instead, which does not do any automatic batching or sharding for you.
For a tutorial on more usage and properties of this method, refer to the tutorial on distributed input. If you are interested in last partial batch handling, read this section.
View source
Returns the list of all local per-replica values contained in value.
View source
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 tf.distribute.Strategy.experimental_distribute_dataset with the returned dataset to further distribute it with the strategy.
View source
Runs ops in fn on each replica, with inputs from input_iterator.
DEPRECATED: This method is not available in TF 2.x. Please switch to using run instead.
When eager execution is enabled, executes ops specified by fn on each replica. Otherwise, builds a graph to execute the ops on each replica.
Each replica will take a single, different input from the inputs provided by one get_next call on the input iterator.
fn may call tf.distribute.get_replica_context() to access members such as replica_id_in_sync_group.
View source
Makes an iterator for input provided via dataset.
DEPRECATED: This method is not available in TF 2.x.
Data from the given dataset will be distributed evenly across all the compute replicas. We will assume that the input dataset is batched by the global batch size. With this assumption, we will make a best effort to divide each batch across all the replicas (one or more workers). If this effort fails, an error will be thrown, and the user should instead use make_input_fn_iterator which provides more control to the user, and does not try to divide a batch across replicas.
The user could also use make_input_fn_iterator if they want to customize which input is fed to which replica/worker etc.
View source
Returns an iterator split across replicas created from an input function.
DEPRECATED: This method is not available in TF 2.x.
The input_fn should take an tf.distribute.InputContext object where information about batching and input sharding can be accessed:
The tf.data.Dataset returned by input_fn should have a per-replica batch size, which may be computed using input_context.get_per_replica_batch_size.
View source
Reduce value across replicas and return result on current device.
To see how this would look with multiple replicas, consider the same example with MirroredStrategy with 2 GPUs:
This API is typically used for aggregating the results returned from different replicas, for reporting etc. For example, loss computed from different replicas can be averaged using this API before printing.
There are a number of different tf.distribute APIs for reducing values across replicas:
What should axis be?
Given a per-replica value returned by run, say a per-example loss, the batch will be divided across all the replicas. This function allows you to aggregate across replicas and optionally also across batch elements by specifying the axis parameter accordingly.
For example, if you have a global batch size of 8 and 2 replicas, values for examples [0, 1, 2, 3] will be on replica 0 and [4, 5, 6, 7] will be on replica 1. With axis=None, reduce will aggregate only across replicas, returning [0+4, 1+5, 2+6, 3+7]. This is useful when each replica is computing a scalar or some other value that doesn't have a "batch" dimension (like a gradient or loss).
Sometimes, you will want to aggregate across both the global batch and all replicas. You can get this behavior by specifying the batch dimension as the axis, typically axis=0. In this case it would return a scalar 0+1+2+3+4+5+6+7.
If there is a last partial batch, you will need to specify an axis so that the resulting shape is consistent across replicas. So if the last batch has size 6 and it is divided into [0, 1, 2, 3] and [4, 5], you would get a shape mismatch unless you specify axis=0. If you specify tf.distribute.ReduceOp.MEAN, using axis=0 will use the correct denominator of 6. Contrast this with computing reduce_mean to get a scalar value on each replica and this function to average those means, which will weigh some values 1/8 and others 1/4.
View source
Invokes fn on each replica, with the given arguments.
This method is the primary way to distribute your computation with a tf.distribute object. It invokes fn on each replica. If args or kwargs have tf.distribute.DistributedValues, such as those produced by a tf.distribute.DistributedDataset from tf.distribute.Strategy.experimental_distribute_dataset or tf.distribute.Strategy.distribute_datasets_from_function, when fn is executed on a particular replica, it will be executed with the component of tf.distribute.DistributedValues that correspond to that replica.
fn is invoked under a replica context. fn may call tf.distribute.get_replica_context() to access members such as all_reduce. Please see the module-level docstring of tf.distribute for the concept of replica context.
All arguments in args or kwargs should either be Python values of a nested structure of tensors, e.g. a list of tensors, in which case args and kwargs will be passed to the fn invoked on each replica. Or args or kwargs can be tf.distribute.DistributedValues containing tensors or composite tensors, i.e. tf.compat.v1.TensorInfo.CompositeTensor, in which case each fn call will get the component of a tf.distribute.DistributedValues corresponding to its replica.
View source
Context manager to make the strategy current and distribute variables.
This method returns a context manager, and is used as follows:
What happens when Strategy.scope is entered?
What should be in scope and what should be outside?
There are a number of requirements on what needs to happen inside the scope. However, in places where we have information about which strategy is in use, we often enter the scope for the user, so they don't have to do it explicitly (i.e. calling those either inside or outside the scope is OK).
View source
Returns a copy of config_proto modified for use with this strategy.
DEPRECATED: This method is not available in TF 2.x.
The updated config has something needed to run a strategy, e.g. configuration to run collective ops, or device filters to improve distributed training performance.


An asynchronous multi-worker parameter server tf.distribute strategy.
Inherits From: Strategy
This strategy requires two roles: workers and parameter servers. Variables and updates to those variables will be assigned to parameter servers and other operations are assigned to workers.
When each worker has more than one GPU, operations will be replicated on all GPUs. Even though operations may be replicated, variables are not and each worker shares a common view for which parameter server a variable is assigned to.
By default it uses TFConfigClusterResolver to detect configurations for multi-worker training. This requires a 'TF_CONFIG' environment variable and the 'TF_CONFIG' must have a cluster spec.
This class assumes each worker is running the same code independently, but parameter servers are running a standard server. This means that while each worker will synchronously compute a single gradient update across all GPUs, updates between workers proceed asynchronously. Operations that occur only on the first replica (such as incrementing the global step), will occur on the first replica of every worker.
It is expected to call call_for_each_replica(fn, ...) for any operations which potentially can be replicated across replicas (i.e. multiple GPUs) even if there is only CPU or one GPU. When defining the fn, extra caution needs to be taken:
1) It is generally not recommended to open a device scope under the strategy's scope. A device scope (i.e. calling tf.device) will be merged with or override the device for operations but will not change the device for variables.
2) It is also not recommended to open a colocation scope (i.e. calling tf.compat.v1.colocate_with) under the strategy's scope. For colocating variables, use strategy.extended.colocate_vars_with instead. Colocation of ops will possibly create device assignment conflicts.
In general, when using a multi-worker tf.distribute strategy such as tf.distribute.experimental.MultiWorkerMirroredStrategy or tf.distribute.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,
For more information, please see tf.distribute.cluster_resolver.ClusterResolver's API docstring.
View source
Distributes tf.data.Dataset instances created by calls to dataset_fn.
The argument dataset_fn that users pass in is an input function that has a tf.distribute.InputContext argument and returns a tf.data.Dataset instance. It is expected that the returned dataset from dataset_fn is already batched by per-replica batch size (i.e. global batch size divided by the number of replicas in sync) and sharded. tf.distribute.Strategy.distribute_datasets_from_function does not batch or shard the tf.data.Dataset instance returned from the input function. dataset_fn will be called on the CPU device of each of the workers and each generates a dataset where every replica on that worker will dequeue one batch of inputs (i.e. if a worker has two replicas, two batches will be dequeued from the Dataset every step).
This method can be used for several purposes. First, it allows you to specify your own batching and sharding logic. (In contrast, tf.distribute.experimental_distribute_dataset does batching and sharding for you.) For example, where experimental_distribute_dataset is unable to shard the input files, this method might be used to manually shard the dataset (avoiding the slow fallback behavior in experimental_distribute_dataset). In cases where the dataset is infinite, this sharding can be done by creating dataset replicas that differ only in their random seed.
The dataset_fn should take an tf.distribute.InputContext instance where information about batching and input replication can be accessed.
You can use element_spec property of the tf.distribute.DistributedDataset returned by this API to query the tf.TypeSpec of the elements returned by the iterator. This can be used to set the input_signature property of a tf.function. Follow tf.distribute.DistributedDataset.element_spec to see an example.
For a tutorial on more usage and properties of this method, refer to the tutorial on distributed input). If you are interested in last partial batch handling, read this section.
View source
Creates tf.distribute.DistributedDataset from tf.data.Dataset.
The returned tf.distribute.DistributedDataset can be iterated over similar to regular datasets. NOTE: The user cannot add any more transformations to a tf.distribute.DistributedDataset. You can only create an iterator or examine the tf.TypeSpec of the data generated by it. See API docs of tf.distribute.DistributedDataset to learn more.
The following is an example:
Three key actions happending under the hood of this method are batching, sharding, and prefetching.
In the code snippet above, dataset is batched by global_batch_size, and calling experimental_distribute_dataset on it rebatches dataset to a new batch size that is equal to the global batch size divided by the number of replicas in sync. We iterate through it using a Pythonic for loop. x is a tf.distribute.DistributedValues containing data for all replicas, and each replica gets data of the new batch size. tf.distribute.Strategy.run will take care of feeding the right per-replica data in x to the right replica_fn executed on each replica.
Sharding contains autosharding across multiple workers and within every worker. First, in multi-worker distributed training (i.e. when you use tf.distribute.experimental.MultiWorkerMirroredStrategy or tf.distribute.TPUStrategy), autosharding a dataset over a set of workers means that each worker is assigned a subset of the entire dataset (if the right tf.data.experimental.AutoShardPolicy is set). This is to ensure that at each step, a global batch size of non-overlapping dataset elements will be processed by each worker. Autosharding has a couple of different options that can be specified using tf.data.experimental.DistributeOptions. Then, sharding within each worker means the method will split the data among all the worker devices (if more than one a present). This will happen regardless of multi-worker autosharding.
By default, this method adds a prefetch transformation at the end of the user provided tf.data.Dataset instance. The argument to the prefetch transformation which is buffer_size is equal to the number of replicas in sync.
If the above batch splitting and dataset sharding logic is undesirable, please use tf.distribute.Strategy.distribute_datasets_from_function instead, which does not do any automatic batching or sharding for you.
For a tutorial on more usage and properties of this method, refer to the tutorial on distributed input. If you are interested in last partial batch handling, read this section.
View source
Returns the list of all local per-replica values contained in value.
View source
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 tf.distribute.Strategy.experimental_distribute_dataset with the returned dataset to further distribute it with the strategy.
View source
Runs ops in fn on each replica, with inputs from input_iterator.
DEPRECATED: This method is not available in TF 2.x. Please switch to using run instead.
When eager execution is enabled, executes ops specified by fn on each replica. Otherwise, builds a graph to execute the ops on each replica.
Each replica will take a single, different input from the inputs provided by one get_next call on the input iterator.
fn may call tf.distribute.get_replica_context() to access members such as replica_id_in_sync_group.
View source
Makes an iterator for input provided via dataset.
DEPRECATED: This method is not available in TF 2.x.
Data from the given dataset will be distributed evenly across all the compute replicas. We will assume that the input dataset is batched by the global batch size. With this assumption, we will make a best effort to divide each batch across all the replicas (one or more workers). If this effort fails, an error will be thrown, and the user should instead use make_input_fn_iterator which provides more control to the user, and does not try to divide a batch across replicas.
The user could also use make_input_fn_iterator if they want to customize which input is fed to which replica/worker etc.
View source
Returns an iterator split across replicas created from an input function.
DEPRECATED: This method is not available in TF 2.x.
The input_fn should take an tf.distribute.InputContext object where information about batching and input sharding can be accessed:
The tf.data.Dataset returned by input_fn should have a per-replica batch size, which may be computed using input_context.get_per_replica_batch_size.
View source
Reduce value across replicas and return result on current device.
To see how this would look with multiple replicas, consider the same example with MirroredStrategy with 2 GPUs:
This API is typically used for aggregating the results returned from different replicas, for reporting etc. For example, loss computed from different replicas can be averaged using this API before printing.
There are a number of different tf.distribute APIs for reducing values across replicas:
What should axis be?
Given a per-replica value returned by run, say a per-example loss, the batch will be divided across all the replicas. This function allows you to aggregate across replicas and optionally also across batch elements by specifying the axis parameter accordingly.
For example, if you have a global batch size of 8 and 2 replicas, values for examples [0, 1, 2, 3] will be on replica 0 and [4, 5, 6, 7] will be on replica 1. With axis=None, reduce will aggregate only across replicas, returning [0+4, 1+5, 2+6, 3+7]. This is useful when each replica is computing a scalar or some other value that doesn't have a "batch" dimension (like a gradient or loss).
Sometimes, you will want to aggregate across both the global batch and all replicas. You can get this behavior by specifying the batch dimension as the axis, typically axis=0. In this case it would return a scalar 0+1+2+3+4+5+6+7.
If there is a last partial batch, you will need to specify an axis so that the resulting shape is consistent across replicas. So if the last batch has size 6 and it is divided into [0, 1, 2, 3] and [4, 5], you would get a shape mismatch unless you specify axis=0. If you specify tf.distribute.ReduceOp.MEAN, using axis=0 will use the correct denominator of 6. Contrast this with computing reduce_mean to get a scalar value on each replica and this function to average those means, which will weigh some values 1/8 and others 1/4.
View source
Invokes fn on each replica, with the given arguments.
This method is the primary way to distribute your computation with a tf.distribute object. It invokes fn on each replica. If args or kwargs have tf.distribute.DistributedValues, such as those produced by a tf.distribute.DistributedDataset from tf.distribute.Strategy.experimental_distribute_dataset or tf.distribute.Strategy.distribute_datasets_from_function, when fn is executed on a particular replica, it will be executed with the component of tf.distribute.DistributedValues that correspond to that replica.
fn is invoked under a replica context. fn may call tf.distribute.get_replica_context() to access members such as all_reduce. Please see the module-level docstring of tf.distribute for the concept of replica context.
All arguments in args or kwargs should either be Python values of a nested structure of tensors, e.g. a list of tensors, in which case args and kwargs will be passed to the fn invoked on each replica. Or args or kwargs can be tf.distribute.DistributedValues containing tensors or composite tensors, i.e. tf.compat.v1.TensorInfo.CompositeTensor, in which case each fn call will get the component of a tf.distribute.DistributedValues corresponding to its replica.
View source
Context manager to make the strategy current and distribute variables.
This method returns a context manager, and is used as follows:
What happens when Strategy.scope is entered?
What should be in scope and what should be outside?
There are a number of requirements on what needs to happen inside the scope. However, in places where we have information about which strategy is in use, we often enter the scope for the user, so they don't have to do it explicitly (i.e. calling those either inside or outside the scope is OK).
View source
Returns a copy of config_proto modified for use with this strategy.
DEPRECATED: This method is not available in TF 2.x.
The updated config has something needed to run a strategy, e.g. configuration to run collective ops, or device filters to improve distributed training performance.