![]() |
![]() |
![]() |
Overview
This tutorial demonstrates how tf.distribute.Strategy
can be used for distributed multi-worker training with tf.estimator
. If you write your code using tf.estimator
, and you're interested in scaling beyond a single machine with high performance, this tutorial is for you.
Before getting started, please read the distribution strategy guide. The multi-GPU training tutorial is also relevant, because this tutorial uses the same model.
Setup
First, setup TensorFlow and the necessary imports.
import tensorflow_datasets as tfds
import tensorflow as tf
import os, json
Input function
This tutorial uses the MNIST dataset from TensorFlow Datasets. The code here is similar to the multi-GPU training tutorial with one key difference: when using Estimator for multi-worker training, it is necessary to shard the dataset by the number of workers to ensure model convergence. The input data is sharded by worker index, so that each worker processes 1/num_workers
distinct portions of the dataset.
BUFFER_SIZE = 10000
BATCH_SIZE = 64
def input_fn(mode, input_context=None):
datasets, info = tfds.load(name='mnist',
with_info=True,
as_supervised=True)
mnist_dataset = (datasets['train'] if mode == tf.estimator.ModeKeys.TRAIN else
datasets['test'])
def scale(image, label):
image = tf.cast(image, tf.float32)
image /= 255
return image, label
if input_context:
mnist_dataset = mnist_dataset.shard(input_context.num_input_pipelines,
input_context.input_pipeline_id)
return mnist_dataset.map(scale).cache().shuffle(BUFFER_SIZE).batch(BATCH_SIZE)
Another reasonable approach to achieve convergence would be to shuffle the dataset with distinct seeds at each worker.
Multi-worker configuration
One of the key differences in this tutorial (compared to the multi-GPU training tutorial) is the multi-worker setup. The TF_CONFIG
environment variable is the standard way to specify the cluster configuration to each worker that is part of the cluster.
There are two components of TF_CONFIG
: cluster
and task
. cluster
provides information about the entire cluster, namely the workers and parameter servers in the cluster. task
provides information about the current task. The first component cluster
is the same for all workers and parameter servers in the cluster, and the second component task
is different on each worker and parameter server and specifies its own type
and index
. In this example, the task type
is worker
and the task index
is 0
.
For illustration purposes, this tutorial shows how to set a TF_CONFIG
with 2 workers on localhost
. In practice, you would create multiple workers on an external IP address and port, and set TF_CONFIG
on each worker appropriately, i.e. modify the task index
.
os.environ['TF_CONFIG'] = json.dumps({
'cluster': {
'worker': ["localhost:12345", "localhost:23456"]
},
'task': {'type': 'worker', 'index': 0}
})
Define the model
Write the layers, the optimizer, and the loss function for training. This tutorial defines the model with Keras layers, similar to the multi-GPU training tutorial.
LEARNING_RATE = 1e-4
def model_fn(features, labels, mode):
model = tf.keras.Sequential([
tf.keras.layers.Conv2D(32, 3, activation='relu', input_shape=(28, 28, 1)),
tf.keras.layers.MaxPooling2D(),
tf.keras.layers.Flatten(),
tf.keras.layers.Dense(64, activation='relu'),
tf.keras.layers.Dense(10)
])
logits = model(features, training=False)
if mode == tf.estimator.ModeKeys.PREDICT:
predictions = {'logits': logits}
return tf.estimator.EstimatorSpec(labels=labels, predictions=predictions)
optimizer = tf.compat.v1.train.GradientDescentOptimizer(
learning_rate=LEARNING_RATE)
loss = tf.keras.losses.SparseCategoricalCrossentropy(
from_logits=True, reduction=tf.keras.losses.Reduction.NONE)(labels, logits)
loss = tf.reduce_sum(loss) * (1. / BATCH_SIZE)
if mode == tf.estimator.ModeKeys.EVAL:
return tf.estimator.EstimatorSpec(mode, loss=loss)
return tf.estimator.EstimatorSpec(
mode=mode,
loss=loss,
train_op=optimizer.minimize(
loss, tf.compat.v1.train.get_or_create_global_step()))
MultiWorkerMirroredStrategy
To train the model, use an instance of tf.distribute.experimental.MultiWorkerMirroredStrategy
. MultiWorkerMirroredStrategy
creates copies of all variables in the model's layers on each device across all workers. It uses CollectiveOps
, a TensorFlow op for collective communication, to aggregate gradients and keep the variables in sync. The tf.distribute.Strategy
guide has more details about this strategy.
strategy = tf.distribute.experimental.MultiWorkerMirroredStrategy()
INFO:tensorflow:Using MirroredStrategy with devices ('/device:GPU:0',) INFO:tensorflow:Single-worker MultiWorkerMirroredStrategy with local_devices = ('/device:GPU:0',), communication = CollectiveCommunication.AUTO
Train and evaluate the model
Next, specify the distribution strategy in the RunConfig
for the estimator, and train and evaluate by invoking tf.estimator.train_and_evaluate
. This tutorial distributes only the training by specifying the strategy via train_distribute
. It is also possible to distribute the evaluation via eval_distribute
.
config = tf.estimator.RunConfig(train_distribute=strategy)
classifier = tf.estimator.Estimator(
model_fn=model_fn, model_dir='/tmp/multiworker', config=config)
tf.estimator.train_and_evaluate(
classifier,
train_spec=tf.estimator.TrainSpec(input_fn=input_fn),
eval_spec=tf.estimator.EvalSpec(input_fn=input_fn)
)
INFO:tensorflow:Initializing RunConfig with distribution strategies. INFO:tensorflow:Not using Distribute Coordinator. INFO:tensorflow:Using config: {'_model_dir': '/tmp/multiworker', '_tf_random_seed': None, '_save_summary_steps': 100, '_save_checkpoints_steps': None, '_save_checkpoints_secs': 600, '_session_config': allow_soft_placement: true graph_options { rewrite_options { meta_optimizer_iterations: ONE } } , '_keep_checkpoint_max': 5, '_keep_checkpoint_every_n_hours': 10000, '_log_step_count_steps': 100, '_train_distribute': <tensorflow.python.distribute.collective_all_reduce_strategy.CollectiveAllReduceStrategy object at 0x7f4c6c18af98>, '_device_fn': None, '_protocol': None, '_eval_distribute': None, '_experimental_distribute': None, '_experimental_max_worker_delay_secs': None, '_session_creation_timeout_secs': 7200, '_service': None, '_cluster_spec': ClusterSpec({}), '_task_type': 'worker', '_task_id': 0, '_global_id_in_cluster': 0, '_master': '', '_evaluation_master': '', '_is_chief': True, '_num_ps_replicas': 0, '_num_worker_replicas': 1, '_distribute_coordinator_mode': None} INFO:tensorflow:Not using Distribute Coordinator. INFO:tensorflow:Running training and evaluation locally (non-distributed). INFO:tensorflow:Start train and evaluate loop. The evaluate will happen after every checkpoint. Checkpoint frequency is determined based on RunConfig arguments: save_checkpoints_steps None or save_checkpoints_secs 600. INFO:tensorflow:The `input_fn` accepts an `input_context` which will be given by DistributionStrategy WARNING:tensorflow:From /tmpfs/src/tf_docs_env/lib/python3.6/site-packages/tensorflow/python/data/ops/multi_device_iterator_ops.py:339: get_next_as_optional (from tensorflow.python.data.ops.iterator_ops) is deprecated and will be removed in a future version. Instructions for updating: Use `tf.data.Iterator.get_next_as_optional()` instead. Warning:tensorflow:From /tmpfs/src/tf_docs_env/lib/python3.6/site-packages/tensorflow/python/data/ops/multi_device_iterator_ops.py:339: get_next_as_optional (from tensorflow.python.data.ops.iterator_ops) is deprecated and will be removed in a future version. Instructions for updating: Use `tf.data.Iterator.get_next_as_optional()` instead. INFO:tensorflow:Calling model_fn. INFO:tensorflow:Calling model_fn. INFO:tensorflow:Done calling model_fn. INFO:tensorflow:Done calling model_fn. Warning:tensorflow:AutoGraph could not transform <function _combine_distributed_scaffold.<locals>.<lambda> at 0x7f4c8dcaf730> and will run it as-is. Cause: could not parse the source code: lambda scaffold: scaffold.ready_op, args=(grouped_scaffold,)) This error may be avoided by creating the lambda in a standalone statement. To silence this warning, decorate the function with @tf.autograph.experimental.do_not_convert Warning:tensorflow:AutoGraph could not transform <function _combine_distributed_scaffold.<locals>.<lambda> at 0x7f4c8dcaf730> and will run it as-is. Cause: could not parse the source code: lambda scaffold: scaffold.ready_op, args=(grouped_scaffold,)) This error may be avoided by creating the lambda in a standalone statement. To silence this warning, decorate the function with @tf.autograph.experimental.do_not_convert Warning: AutoGraph could not transform <function _combine_distributed_scaffold.<locals>.<lambda> at 0x7f4c8dcaf730> and will run it as-is. Cause: could not parse the source code: lambda scaffold: scaffold.ready_op, args=(grouped_scaffold,)) This error may be avoided by creating the lambda in a standalone statement. To silence this warning, decorate the function with @tf.autograph.experimental.do_not_convert INFO:tensorflow:Create CheckpointSaverHook. INFO:tensorflow:Create CheckpointSaverHook. Warning:tensorflow:From /tmpfs/src/tf_docs_env/lib/python3.6/site-packages/tensorflow_estimator/python/estimator/util.py:96: DistributedIteratorV1.initialize (from tensorflow.python.distribute.input_lib) is deprecated and will be removed in a future version. Instructions for updating: Use the iterator's `initializer` property instead. Warning:tensorflow:From /tmpfs/src/tf_docs_env/lib/python3.6/site-packages/tensorflow_estimator/python/estimator/util.py:96: DistributedIteratorV1.initialize (from tensorflow.python.distribute.input_lib) is deprecated and will be removed in a future version. Instructions for updating: Use the iterator's `initializer` property instead. INFO:tensorflow:Graph was finalized. INFO:tensorflow:Graph was finalized. INFO:tensorflow:Running local_init_op. INFO:tensorflow:Running local_init_op. INFO:tensorflow:Done running local_init_op. INFO:tensorflow:Done running local_init_op. INFO:tensorflow:Calling checkpoint listeners before saving checkpoint 0... INFO:tensorflow:Calling checkpoint listeners before saving checkpoint 0... INFO:tensorflow:Saving checkpoints for 0 into /tmp/multiworker/model.ckpt. INFO:tensorflow:Saving checkpoints for 0 into /tmp/multiworker/model.ckpt. INFO:tensorflow:Calling checkpoint listeners after saving checkpoint 0... INFO:tensorflow:Calling checkpoint listeners after saving checkpoint 0... INFO:tensorflow:loss = 2.3278575, step = 0 INFO:tensorflow:loss = 2.3278575, step = 0 INFO:tensorflow:global_step/sec: 201.897 INFO:tensorflow:global_step/sec: 201.897 INFO:tensorflow:loss = 2.3006024, step = 100 (0.498 sec) INFO:tensorflow:loss = 2.3006024, step = 100 (0.498 sec) INFO:tensorflow:global_step/sec: 215.773 INFO:tensorflow:global_step/sec: 215.773 INFO:tensorflow:loss = 2.2919793, step = 200 (0.463 sec) INFO:tensorflow:loss = 2.2919793, step = 200 (0.463 sec) INFO:tensorflow:global_step/sec: 213.717 INFO:tensorflow:global_step/sec: 213.717 INFO:tensorflow:loss = 2.286222, step = 300 (0.468 sec) INFO:tensorflow:loss = 2.286222, step = 300 (0.468 sec) INFO:tensorflow:global_step/sec: 215.652 INFO:tensorflow:global_step/sec: 215.652 INFO:tensorflow:loss = 2.2875795, step = 400 (0.464 sec) INFO:tensorflow:loss = 2.2875795, step = 400 (0.464 sec) INFO:tensorflow:global_step/sec: 215.686 INFO:tensorflow:global_step/sec: 215.686 INFO:tensorflow:loss = 2.3000607, step = 500 (0.466 sec) INFO:tensorflow:loss = 2.3000607, step = 500 (0.466 sec) INFO:tensorflow:global_step/sec: 217.858 INFO:tensorflow:global_step/sec: 217.858 INFO:tensorflow:loss = 2.2862964, step = 600 (0.457 sec) INFO:tensorflow:loss = 2.2862964, step = 600 (0.457 sec) INFO:tensorflow:global_step/sec: 216.886 INFO:tensorflow:global_step/sec: 216.886 INFO:tensorflow:loss = 2.2848775, step = 700 (0.463 sec) INFO:tensorflow:loss = 2.2848775, step = 700 (0.463 sec) INFO:tensorflow:global_step/sec: 242.69 INFO:tensorflow:global_step/sec: 242.69 INFO:tensorflow:loss = 2.2776775, step = 800 (0.409 sec) INFO:tensorflow:loss = 2.2776775, step = 800 (0.409 sec) INFO:tensorflow:global_step/sec: 621.93 INFO:tensorflow:global_step/sec: 621.93 INFO:tensorflow:loss = 2.283049, step = 900 (0.161 sec) INFO:tensorflow:loss = 2.283049, step = 900 (0.161 sec) INFO:tensorflow:Calling checkpoint listeners before saving checkpoint 938... INFO:tensorflow:Calling checkpoint listeners before saving checkpoint 938... INFO:tensorflow:Saving checkpoints for 938 into /tmp/multiworker/model.ckpt. INFO:tensorflow:Saving checkpoints for 938 into /tmp/multiworker/model.ckpt. INFO:tensorflow:Calling checkpoint listeners after saving checkpoint 938... INFO:tensorflow:Calling checkpoint listeners after saving checkpoint 938... INFO:tensorflow:Calling model_fn. INFO:tensorflow:Calling model_fn. INFO:tensorflow:Done calling model_fn. INFO:tensorflow:Done calling model_fn. INFO:tensorflow:Starting evaluation at 2020-09-11T01:27:54Z INFO:tensorflow:Starting evaluation at 2020-09-11T01:27:54Z INFO:tensorflow:Graph was finalized. INFO:tensorflow:Graph was finalized. INFO:tensorflow:Restoring parameters from /tmp/multiworker/model.ckpt-938 INFO:tensorflow:Restoring parameters from /tmp/multiworker/model.ckpt-938 INFO:tensorflow:Running local_init_op. INFO:tensorflow:Running local_init_op. INFO:tensorflow:Done running local_init_op. INFO:tensorflow:Done running local_init_op. INFO:tensorflow:Evaluation [10/100] INFO:tensorflow:Evaluation [10/100] INFO:tensorflow:Evaluation [20/100] INFO:tensorflow:Evaluation [20/100] INFO:tensorflow:Evaluation [30/100] INFO:tensorflow:Evaluation [30/100] INFO:tensorflow:Evaluation [40/100] INFO:tensorflow:Evaluation [40/100] INFO:tensorflow:Evaluation [50/100] INFO:tensorflow:Evaluation [50/100] INFO:tensorflow:Evaluation [60/100] INFO:tensorflow:Evaluation [60/100] INFO:tensorflow:Evaluation [70/100] INFO:tensorflow:Evaluation [70/100] INFO:tensorflow:Evaluation [80/100] INFO:tensorflow:Evaluation [80/100] INFO:tensorflow:Evaluation [90/100] INFO:tensorflow:Evaluation [90/100] INFO:tensorflow:Evaluation [100/100] INFO:tensorflow:Evaluation [100/100] INFO:tensorflow:Inference Time : 1.01975s INFO:tensorflow:Inference Time : 1.01975s INFO:tensorflow:Finished evaluation at 2020-09-11-01:27:55 INFO:tensorflow:Finished evaluation at 2020-09-11-01:27:55 INFO:tensorflow:Saving dict for global step 938: global_step = 938, loss = 2.276255 INFO:tensorflow:Saving dict for global step 938: global_step = 938, loss = 2.276255 INFO:tensorflow:Saving 'checkpoint_path' summary for global step 938: /tmp/multiworker/model.ckpt-938 INFO:tensorflow:Saving 'checkpoint_path' summary for global step 938: /tmp/multiworker/model.ckpt-938 INFO:tensorflow:Loss for final step: 1.1389045. INFO:tensorflow:Loss for final step: 1.1389045. ({'loss': 2.276255, 'global_step': 938}, [])
Optimize training performance
You now have a model and a multi-worker capable Estimator powered by tf.distribute.Strategy
. You can try the following techniques to optimize performance of multi-worker training:
- Increase the batch size: The batch size specified here is per-GPU. In general, the largest batch size that fits the GPU memory is advisable.
- Cast variables: Cast the variables to
tf.float
if possible. The official ResNet model includes an example of how this can be done. Use collective communication:
MultiWorkerMirroredStrategy
provides multiple collective communication implementations.RING
implements ring-based collectives using gRPC as the cross-host communication layer.NCCL
uses Nvidia's NCCL to implement collectives.AUTO
defers the choice to the runtime.
The best choice of collective implementation depends upon the number and kind of GPUs, and the network interconnect in the cluster. To override the automatic choice, specify a valid value to the
communication
parameter ofMultiWorkerMirroredStrategy
's constructor, e.g.communication=tf.distribute.experimental.CollectiveCommunication.NCCL
.
Visit the Performance section in the guide to learn more about other strategies and tools you can use to optimize the performance of your TensorFlow models.
Other code examples
- End to end example for multi worker training in tensorflow/ecosystem using Kubernetes templates. This example starts with a Keras model and converts it to an Estimator using the
tf.keras.estimator.model_to_estimator
API. - Official models, many of which can be configured to run multiple distribution strategies.