![]() |
TPU distribution strategy implementation.
Inherits From: Strategy
tf.compat.v1.distribute.experimental.TPUStrategy(
tpu_cluster_resolver=None, steps_per_run=None, device_assignment=None
)
Args | |
---|---|
tpu_cluster_resolver
|
A tf.distribute.cluster_resolver.TPUClusterResolver, which provides information about the TPU cluster. |
steps_per_run
|
Number of steps to run on device before returning to the host. Note that this can have side-effects on performance, hooks, metrics, summaries etc. This parameter is only used when Distribution Strategy is used with estimator or keras. |
device_assignment
|
Optional tf.tpu.experimental.DeviceAssignment to
specify the placement of replicas on the TPU cluster. Currently only
supports the usecase of using a single core within a TPU cluster.
|
Attributes | |
---|---|
cluster_resolver
|
Returns the cluster resolver associated with this strategy.
In general, when using a multi-worker Strategies that intend to have an associated
Single-worker strategies usually do not have a
The
For more information, please see
|
extended
|
tf.distribute.StrategyExtended with additional methods.
|
num_replicas_in_sync
|
Returns number of replicas over which gradients are aggregated. |
steps_per_run
|
DEPRECATED: use .extended.steps_per_run instead. |
Methods
distribute_datasets_from_function
distribute_datasets_from_function(
dataset_fn, options=None
)
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.
Args | |
---|---|
dataset_fn
|
A function taking a tf.distribute.InputContext instance and
returning a tf.data.Dataset .
|
options
|
tf.distribute.InputOptions used to control options on how this
dataset is distributed.
|
Returns | |
---|---|
A tf.distribute.DistributedDataset .
|
experimental_distribute_dataset
experimental_distribute_dataset(
dataset, options=None
)
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:
global_batch_size = 2
# Passing the devices is optional.
strategy = tf.distribute.MirroredStrategy(devices=["GPU:0", "GPU:1"])
# Create a dataset
dataset = tf.data.Dataset.range(4).batch(global_batch_size)
# Distribute that dataset
dist_dataset = strategy.experimental_distribute_dataset(dataset)
@tf.function
def replica_fn(input):
return input*2
result = []
# Iterate over the `tf.distribute.DistributedDataset`
for x in dist_dataset:
# process dataset elements
result.append(strategy.run(replica_fn, args=(x,)))
print(result)
[PerReplica:{
0: <tf.Tensor: shape=(1,), dtype=int64, numpy=array([0])>,
1: <tf.Tensor: shape=(1,), dtype=int64, numpy=array([2])>
}, PerReplica:{
0: <tf.Tensor: shape=(1,), dtype=int64, numpy=array([4])>,
1: <tf.Tensor: shape=(1,), dtype=int64, numpy=array([6])>
}]
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.
Args | |
---|---|
dataset
|
tf.data.Dataset that will be sharded across all replicas using
the rules stated above.
|
options
|
tf.distribute.InputOptions used to control options on how this
dataset is distributed.
|
Returns | |
---|---|
A tf.distribute.DistributedDataset .
|
experimental_local_results
experimental_local_results(
value
)
Returns the list of all local per-replica values contained in value
.
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
experimental_make_numpy_dataset(
numpy_input, session=None
)
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.
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.
|
session
|
(TensorFlow v1.x graph execution only) A session used for initialization. |
Returns | |
---|---|
A tf.data.Dataset representing numpy_input .
|
experimental_run
experimental_run(
fn, input_iterator=None
)
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
.
Args | |
---|---|
fn
|
The function to run. The inputs to the function must match the outputs
of input_iterator.get_next() . The output must be a tf.nest of
Tensor s.
|
input_iterator
|
(Optional) input iterator from which the inputs are taken. |
Returns | |
---|---|
Merged return value of fn across replicas. The structure of the return
value is the same as the return value from fn . Each element in the
structure can either be PerReplica (if the values are unsynchronized),
Mirrored (if the values are kept in sync), or Tensor (if running on a
single replica).
|
make_dataset_iterator
make_dataset_iterator(
dataset
)
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.
Args | |
---|---|
dataset
|