|View source on GitHub|
Synchronous training on TPUs and TPU Pods.
tf.distribute.TPUStrategy( tpu_cluster_resolver=None, experimental_device_assignment=None )
Used in the notebooks
|Used in the guide|
To construct a TPUStrategy object, you need to run the initialization code as below:
resolver = tf.distribute.cluster_resolver.TPUClusterResolver(tpu='')
strategy = tf.distribute.TPUStrategy(resolver)
While using distribution strategies, the variables created within the strategy's scope will be replicated across all the replicas and can be kept in sync using all-reduce algorithms.
To run TF2 programs on TPUs, you can either use
.fit APIs in
tf.keras with TPUStrategy, or write your own customized
training loop by calling
strategy.run directly. Note that
TPUStrategy doesn't support pure eager execution, so please make sure the
function passed into
strategy.run is a
strategy.run is called inside a
tf.function if eager
behavior is enabled. See more details in https://www.tensorflow.org/guide/tpu.
experimental_distribute_dataset APIs can be used to distribute the dataset
across the TPU workers when writing your own training loop. If you are using
compile methods available in
tf.keras.Model, then Keras will
handle the distribution for you.
An example of writing customized training loop on TPUs:
model = tf.keras.Sequential([
optimizer = tf.keras.optimizers.SGD(learning_rate=0.1)
x = np.random.random((2, 5)).astype(np.float32)