|---|---|---|---| | [![](https://www.tensorflow.org/images/tf_logo_32px.png)View on TensorFlow.org](https://www.tensorflow.org/tutorials/distribute/multi_worker_with_keras) | [![](https://www.tensorflow.org/images/colab_logo_32px.png)Run in Google Colab](https://colab.research.google.com/github/tensorflow/docs/blob/master/site/en/tutorials/distribute/multi_worker_with_keras.ipynb) | [![](https://www.tensorflow.org/images/GitHub-Mark-32px.png)View source on GitHub](https://github.com/tensorflow/docs/blob/master/site/en/tutorials/distribute/multi_worker_with_keras.ipynb) | [![](https://www.tensorflow.org/images/download_logo_32px.png)Download notebook](https://storage.googleapis.com/tensorflow_docs/docs/site/en/tutorials/distribute/multi_worker_with_keras.ipynb) | ## Overview This tutorial demonstrates how to perform multi-worker distributed training with a Keras model and the [`Model.fit`](https://www.tensorflow.org/api_docs/python/tf/keras/Model#fit) API using the [`tf.distribute.MultiWorkerMirroredStrategy`](https://www.tensorflow.org/api_docs/python/tf/distribute/MultiWorkerMirroredStrategy) API. With the help of this strategy, a Keras model that was designed to run on a single-worker can seamlessly work on multiple workers with minimal code changes. To learn how to use the `MultiWorkerMirroredStrategy` with Keras and a custom training loop, refer to [Custom training loop with Keras and MultiWorkerMirroredStrategy](https://www.tensorflow.org/tutorials/distribute/multi_worker_with_ctl). This tutorial contains a minimal multi-worker example with two workers for demonstration purposes. ### Choose the right strategy Before you dive in, make sure that [`tf.distribute.MultiWorkerMirroredStrategy`](https://www.tensorflow.org/api_docs/python/tf/distribute/MultiWorkerMirroredStrategy) is the right choice for your accelerator(s) and training. These are two common ways of distributing training with data parallelism: - *Synchronous training* , where the steps of training are synced across the workers and replicas, such as [`tf.distribute.MirroredStrategy`](https://www.tensorflow.org/api_docs/python/tf/distribute/MirroredStrategy), [`tf.distribute.TPUStrategy`](https://www.tensorflow.org/api_docs/python/tf/distribute/TPUStrategy), and [`tf.distribute.MultiWorkerMirroredStrategy`](https://www.tensorflow.org/api_docs/python/tf/distribute/MultiWorkerMirroredStrategy). All workers train over different slices of input data in sync, and aggregating gradients at each step. - *Asynchronous training* , where the training steps are not strictly synced, such as [`tf.distribute.experimental.ParameterServerStrategy`](https://www.tensorflow.org/api_docs/python/tf/distribute/experimental/ParameterServerStrategy). All workers are independently training over the input data and updating variables asynchronously. If you are looking for multi-worker synchronous training without TPU, then [`tf.distribute.MultiWorkerMirroredStrategy`](https://www.tensorflow.org/api_docs/python/tf/distribute/MultiWorkerMirroredStrategy) is your choice. It 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 keeps the variables in sync. For those interested, check out the [`tf.distribute.experimental.CommunicationOptions`](https://www.tensorflow.org/api_docs/python/tf/distribute/experimental/CommunicationOptions) parameter for the collective implementation options. For an overview of [`tf.distribute.Strategy`](https://www.tensorflow.org/api_docs/python/tf/distribute/Strategy) APIs, refer to [Distributed training in TensorFlow](https://www.tensorflow.org/guide/distributed_training). ## Setup Start with some necessary imports: import json import os import sys Before importing TensorFlow, make a few changes to the environment: - In a real-world application, each worker would be on a different machine. For the purposes of this tutorial, all the workers will run on the **this** machine. Therefore, disable all GPUs to prevent errors caused by all workers trying to use the same GPU. os.environ["CUDA_VISIBLE_DEVICES"] = "-1" - Reset the `TF_CONFIG` environment variable (you'll learn more about this later): os.environ.pop('TF_CONFIG', None) - Make sure that the current directory is on Python's path---this allows the notebook to import the files written by `%%writefile` later: if '.' not in sys.path: sys.path.insert(0, '.') Install `tf-nightly`, as the frequency of checkpoint saving at a particular step with the `save_freq` argument in `tf.keras.callbacks.BackupAndRestore` is introduced from TensorFlow 2.10: pip install tf-nightly Finally, import TensorFlow: import tensorflow as tf ### Dataset and model definition Next, create an `mnist_setup.py` file with a simple model and dataset setup. This Python file will be used by the worker processes in this tutorial: %%writefile mnist_setup.py import os import tensorflow as tf import numpy as np def mnist_dataset(batch_size): (x_train, y_train), _ = tf.keras.datasets.mnist.load_data() # The `x` arrays are in uint8 and have values in the [0, 255] range. # You need to convert them to float32 with values in the [0, 1] range. x_train = x_train / np.float32(255) y_train = y_train.astype(np.int64) train_dataset = tf.data.Dataset.from_tensor_slices( (x_train, y_train)).shuffle(60000).repeat().batch(batch_size) return train_dataset def build_and_compile_cnn_model(): model = tf.keras.Sequential([ tf.keras.layers.InputLayer(input_shape=(28, 28)), tf.keras.layers.Reshape(target_shape=(28, 28, 1)), tf.keras.layers.Conv2D(32, 3, activation='relu'), tf.keras.layers.Flatten(), tf.keras.layers.Dense(128, activation='relu'), tf.keras.layers.Dense(10) ]) model.compile( loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True), optimizer=tf.keras.optimizers.SGD(learning_rate=0.001), metrics=['accuracy']) return model ``` Writing mnist_setup.py ``` ### Model training on a single worker Try training the model for a small number of epochs and observe the results of *a single worker* to make sure everything works correctly. As training progresses, the loss should drop and the accuracy should increase. import mnist_setup batch_size = 64 single_worker_dataset = mnist_setup.mnist_dataset(batch_size) single_worker_model = mnist_setup.build_and_compile_cnn_model() single_worker_model.fit(single_worker_dataset, epochs=3, steps_per_epoch=70) ``` WARNING: All log messages before absl::InitializeLog() is called are written to STDERR E0000 00:00:1793850180.824325 80694 cuda_platform.cc:57] failed call to cuInit: INTERNAL: CUDA error: Failed call to cuInit: CUDA_ERROR_NO_DEVICE: no CUDA-capable device is detected === Source Location Trace: === external/xla/xla/stream_executor/cuda/cuda_status.cc:50 Epoch 1/3 70/70 [==============================] - 1s 7ms/step - loss: 2.2819 - accuracy: 0.1580 Epoch 2/3 70/70 [==============================] - 0s 6ms/step - loss: 2.2475 - accuracy: 0.3089 Epoch 3/3 70/70 [==============================] - 0s 7ms/step - loss: 2.2011 - accuracy: 0.4583 ``` ## Multi-worker configuration Now let's enter the world of multi-worker training. ### A cluster with jobs and tasks In TensorFlow, distributed training involves a `'cluster'` with several jobs, and each of the jobs may have one or more `'task'`s. You will need the `TF_CONFIG` configuration environment variable for training on multiple machines, each of which possibly has a different role. `TF_CONFIG` is a JSON string used to specify the cluster configuration for each worker that is part of the cluster. There are two components of a `TF_CONFIG` variable: `'cluster'` and `'task'`. - A `'cluster'` is the same for all workers and provides information about the training cluster, which is a dict consisting of different types of jobs, such as `'worker'` or `'chief'`. - In multi-worker training with [`tf.distribute.MultiWorkerMirroredStrategy`](https://www.tensorflow.org/api_docs/python/tf/distribute/MultiWorkerMirroredStrategy), there is usually one `'worker'` that takes on more responsibilities, such as saving a checkpoint and writing a summary file for TensorBoard, in addition to what a regular `'worker'` does. Such `'worker'` is referred to as the chief worker (with a job name `'chief'`). - It is customary for the worker with `'index'` `0` to be the `'chief'`. - A `'task'` provides information on the current task and is different for each worker. It specifies the `'type'` and `'index'` of that worker. Below is an example configuration: tf_config = { 'cluster': { 'worker': ['localhost:12345', 'localhost:23456'] }, 'task': {'type': 'worker', 'index': 0} } Note that `tf_config` is just a local variable in Python. To use it for training configuration, serialize it as a JSON and place it in a `TF_CONFIG` environment variable. json.dumps(tf_config) ``` '{"cluster": {"worker": ["localhost:12345", "localhost:23456"]}, "task": {"type": "worker", "index": 0} }' ``` In the example configuration above, you set the task `'type'` to `'worker'` and the task `'index'` to `0`. Therefore, this machine is the *first* worker. It will be appointed as the `'chief'` worker. > [!NOTE] > **Note:** Other machines will need to have the `TF_CONFIG` environment variable set as well, and it should have the same `'cluster'` dict, but different task `'type'`s or task `'index'`es, depending on the roles of those machines. In practice, you would create multiple workers on external IP addresses/ports and set a `TF_CONFIG` variable on each worker accordingly. For illustration purposes, this tutorial shows how you may set up a `TF_CONFIG` variable with two workers on a `localhost`: - The first (`'chief'`) worker's `TF_CONFIG` as shown above. - For the second worker, you will set `tf_config['task']['index']=1` ### Environment variables and subprocesses in notebooks Subprocesses inherit environment variables from their parent. So if you set an environment variable in this Jupyter Notebook process: os.environ['GREETINGS'] = 'Hello TensorFlow!' ... then you can access the environment variable from the subprocesses: echo ${GREETINGS} ``` Hello TensorFlow! ``` In the next section, you'll use this method to pass the `TF_CONFIG` to the worker subprocesses. You would never really launch your jobs this way in a real-world scenario---this tutorial is just showing how to do it with a minimal multi-worker example. ## Train the model To train the model, firstly create an instance of the [`tf.distribute.MultiWorkerMirroredStrategy`](https://www.tensorflow.org/api_docs/python/tf/distribute/MultiWorkerMirroredStrategy): strategy = tf.distribute.MultiWorkerMirroredStrategy() ``` WARNING:tensorflow:Collective ops is not configured at program startup. Some performance features may not be enabled. INFO:tensorflow:Using MirroredStrategy with devices ('/device:CPU:0',) INFO:tensorflow:Single-worker MultiWorkerMirroredStrategy with local_devices = ('/device:CPU:0',), communication = CommunicationImplementation.AUTO ``` > [!NOTE] > **Note:** `TF_CONFIG` is parsed and TensorFlow's GRPC servers are started at the time `MultiWorkerMirroredStrategy` is called, so the `TF_CONFIG` environment variable must be set before a [`tf.distribute.Strategy`](https://www.tensorflow.org/api_docs/python/tf/distribute/Strategy) instance is created. Since `TF_CONFIG` is not set yet, the above strategy is effectively single-worker training. With the integration of [`tf.distribute.Strategy`](https://www.tensorflow.org/api_docs/python/tf/distribute/Strategy) API into [`tf.keras`](https://www.tensorflow.org/api_docs/python/tf/keras), the only change you will make to distribute the training to multiple-workers is enclosing the model building and `model.compile()` call inside `strategy.scope()`. The distribution strategy's scope dictates how and where the variables are created, and in the case of `MultiWorkerMirroredStrategy`, the variables created are `MirroredVariable`s, and they are replicated on each of the workers. with strategy.scope(): # Model building/compiling need to be within `strategy.scope()`. multi_worker_model = mnist_setup.build_and_compile_cnn_model() > [!NOTE] > **Note:** Currently there is a limitation in `MultiWorkerMirroredStrategy` where TensorFlow ops need to be created after the instance of strategy is created. If you encounter `RuntimeError: Collective ops must be configured at program startup`, try creating the instance of `MultiWorkerMirroredStrategy` at the beginning of the program and put the code that may create ops after the strategy is instantiated. To actually run with `MultiWorkerMirroredStrategy` you'll need to run worker processes and pass a `TF_CONFIG` to them. Like the `mnist_setup.py` file written earlier, here is the `main.py` that each of the workers will run: %%writefile main.py import os import json import tensorflow as tf import mnist_setup per_worker_batch_size = 64 tf_config = json.loads(os.environ['TF_CONFIG']) num_workers = len(tf_config['cluster']['worker']) strategy = tf.distribute.MultiWorkerMirroredStrategy() global_batch_size = per_worker_batch_size * num_workers multi_worker_dataset = mnist_setup.mnist_dataset(global_batch_size) with strategy.scope(): # Model building/compiling need to be within `strategy.scope()`. multi_worker_model = mnist_setup.build_and_compile_cnn_model() multi_worker_model.fit(multi_worker_dataset, epochs=3, steps_per_epoch=70) ``` Writing main.py ``` In the code snippet above note that the `global_batch_size`, which gets passed to [`Dataset.batch`](https://www.tensorflow.org/api_docs/python/tf/data/Dataset#batch), is set to `per_worker_batch_size * num_workers`. This ensures that each worker processes batches of `per_worker_batch_size` examples regardless of the number of workers. The current directory now contains both Python files: ls *.py ``` main.py mnist_setup.py ``` Serialize the `TF_CONFIG` to JSON and add it to the environment variables: os.environ['TF_CONFIG'] = json.dumps(tf_config) Now, you can launch a worker process that will run the `main.py` and use the `TF_CONFIG`: # first kill any previous runs %killbgscripts ``` All background processes were killed. ``` python main.py &> job_0.log There are a few things to note about the above command: 1. It uses the `%%bash` which is a [notebook "magic"](https://ipython.readthedocs.io/en/stable/interactive/magics.html) to run some bash commands. 2. It uses the `--bg` flag to run the `bash` process in the background, because this worker will not terminate. It waits for all the workers before it starts. The backgrounded worker process won't print output to this notebook, so the `&>` redirects its output to a file so that you can inspect what happened in a log file later. So, wait a few seconds for the process to start up: import time time.sleep(10) Now, inspect what's been output to the worker's log file so far: cat job_0.log ``` Error processing line 1 of /tmpfs/src/tf_docs_env/lib/python3.11/site-packages/tfdocs_sitecustomize. pth: Traceback (most recent call last): File "", line 195, in addpackage Fil e "", line 1 import _tfdocs_sitecustomize\n ^ Syn taxError: unexpected character after line continuation character Remainder of file ignored Error pr ocessing line 1 of /tmpfs/src/tf_docs_env/lib/python3.11/site-packages/tfdocs_sitecustomize.pth: Traceback (most recent call last): File "", line 195, in addpackage File "", line 1 import _tfdocs_sitecustomize\n ^ SyntaxError : unexpected character after line continuation character Remainder of file ignored WARNING: All log messages before absl::InitializeLog() is called are written to STDERR E0000 00:00:1793850186.812013 80889 cuda_platform.cc:57] failed call to cuInit: INTERNAL: CUDA error: Failed call to cuInit: CU DA_ERROR_NO_DEVICE: no CUDA-capable device is detected === Source Location Trace: === external/xla/ xla/stream_executor/cuda/cuda_status.cc:50 ``` The last line of the log file should say: `Started server with target: grpc://localhost:12345`. The first worker is now ready and is waiting for all the other worker(s) to be ready to proceed. So update the `tf_config` for the second worker's process to pick up: tf_config['task']['index'] = 1 os.environ['TF_CONFIG'] = json.dumps(tf_config) Launch the second worker. This will start the training since all the workers are active (so there's no need to background this process): python main.py ``` Error processing line 1 of /tmpfs/src/tf_docs_env/lib/python3.11/site-packages/tfdocs_sitecustomize. pth: Traceback (most recent call last): File "", line 195, in addpackage File "", line 1 import _tfdocs_sitecustomize\n ^ SyntaxErr or: unexpected character after line continuation character Remainder of file ignored Error processing line 1 of /tmpfs/src/tf_docs_env/lib/python3.11/site-packages/tfdocs_sitecustomize. pth: Traceback (most recent call last): File "", line 195, in addpackage Fil e "", line 1 import _tfdocs_sitecustomize\n ^ Syn taxError: unexpected character after line continuation character Remainder of file ignored WARNING: All log messages before absl::InitializeLog() is called are written to STDERR E0000 00:00:1 793850196.917130 81201 cuda_platform.cc:57] failed call to cuInit: INTERNAL: CUDA error: Failed ca ll to cuInit: CUDA_ERROR_NO_DEVICE: no CUDA-capable device is detected === Source Location Trace: == = external/xla/xla/stream_executor/cuda/cuda_status.cc:50 W0000 00:00:1793850198.184790 81201 dataset.cc:1000] Input of GeneratorDatasetOp::Dataset will not be optimized because the dataset does not implement the AsGraphDefInternal() method needed to apply optimizations. Epoch 1/3 2/70 [......... 3/70 [>......... 4/70 [>......... 5/70 [=>........ 6/70 [=>........ 7/70 [==>....... 8/70 [==>....... 9/70 [==>....... 10/70 [===>...... 11/70 [===>...... 12/70 [====>..... 13/70 [====>..... 14/70 [=====>.... 15/70 [=====>.... 16/70 [=====>.... 17/70 [======>... 18/70 [======>... 19/70 [=======>.. 20/70 [=======>.. 21/70 [========>. 22/70 [========>. 23/70 [========>. 24/70 [=========> 25/70 [=========> 26/70 [========== 27/70 [========== 28/70 [========== 29/70 [========== 30/70 [========== 31/70 [========== 32/70 [========== 33/70 [========== 34/70 [========== 35/70 [========== 36/70 [========== 37/70 [========== 38/70 [========== 39/70 [========== 40/70 [========== 41/70 [========== 42/70 [========== 43/70 [========== 44/70 [========== 45/70 [========== 46/70 [========== 47/70 [========== 48/70 [========== 49/70 [========== 50/70 [========== 51/70 [========== 52/70 [========== 53/70 [========== 54/70 [========== 55/70 [========== 56/70 [========== 57/70 [========== 58/70 [========== 59/70 [========== 60/70 [========== 61/70 [========== 62/70 [========== 63/70 [========== 64/70 [========== 65/70 [========== 66/70 [========== 67/70 [========== 68/70 [========== 69/70 [========== 70/70 [========== 70/70 [========== ====================] - 6s 70ms/step - loss: 2.2850 - accuracy: 0.1368 Epoch 2/3 2/70 [.......... 3/70 [>......... 4/70 [>......... 5/70 [=>........ 6/70 [=>........ 7/70 [==>....... 8/70 [==>....... 9/70 [==>....... 10/70 [===>...... 11/70 [===>...... 12/70 [====>..... 13/70 [====>..... 14/70 [=====>.... 15/70 [=====>.... 16/70 [=====>.... 17/70 [======>... 18/70 [======>... 19/70 [=======>.. 20/70 [=======>.. 21/70 [========>. 22/70 [========>. 23/70 [========>. 24/70 [=========> 25/70 [=========> 26/70 [========== 27/70 [========== 28/70 [========== 29/70 [========== 30/70 [========== 31/70 [========== 32/70 [========== 33/70 [========== 34/70 [========== 35/70 [========== 36/70 [========== 37/70 [========== 38/70 [========== 39/70 [========== 40/70 [========== 41/70 [========== 42/70 [========== 43/70 [========== 44/70 [========== 45/70 [========== 46/70 [========== 47/70 [========== 48/70 [========== 49/70 [========== 50/70 [========== 51/70 [========== 52/70 [========== 53/70 [========== 54/70 [========== 55/70 [========== 56/70 [========== 57/70 [========== 58/70 [========== 59/70 [========== 60/70 [========== 61/70 [========== 62/70 [========== 63/70 [========== 64/70 [========== 65/70 [========== 66/70 [========== 67/70 [========== 68/70 [========== 69/70 [========== 70/70 [========== 70/70 [========== ====================] - 5s 67ms/step - loss: 2.2243 - accuracy: 0.3095 Epoch 3/3 2/70 [.......... 3/70 [>......... 4/70 [>......... 5/70 [=>........ 6/70 [=>........ 7/70 [==>....... 8/70 [==>....... 9/70 [==>....... 10/70 [===>...... 11/70 [===>...... 12/70 [====>..... 13/70 [====>..... 14/70 [=====>.... 15/70 [=====>.... 16/70 [=====>.... 17/70 [======>... 18/70 [======>... 19/70 [=======>.. 20/70 [=======>.. 21/70 [========>. 22/70 [========>. 23/70 [========>. 24/70 [=========> 25/70 [=========> 26/70 [========== 27/70 [========== 28/70 [========== 29/70 [========== 30/70 [========== 31/70 [========== 32/70 [========== 33/70 [========== 34/70 [========== 35/70 [========== 36/70 [========== 37/70 [========== 38/70 [========== 39/70 [========== 40/70 [========== 41/70 [========== 42/70 [========== 43/70 [========== 44/70 [========== 45/70 [========== 46/70 [========== 47/70 [========== 48/70 [========== 49/70 [========== 50/70 [========== 51/70 [========== 52/70 [========== 53/70 [========== 54/70 [========== 55/70 [========== 56/70 [========== 57/70 [========== 58/70 [========== 59/70 [========== 60/70 [========== 61/70 [========== 62/70 [========== 63/70 [========== 64/70 [========== 65/70 [========== 66/70 [========== 67/70 [========== 68/70 [========== 69/70 [========== 70/70 [========== 70/70 [========== ====================] - 5s 67ms/step - loss: 2.1596 - accuracy: 0.4751 ``` If you recheck the logs written by the first worker, you'll learn that it participated in training that model: cat job_0.log ``` Error processing line 1 of /tmpfs/src/tf_docs_env/lib/python3.11/site-packages/tfdocs_sitecustomize. pth: Traceback (most recent call last): File "", line 195, in addpackage Fil e "", line 1 import _tfdocs_sitecustomize\n ^ Syn taxError: unexpected character after line continuation character Remainder of file ignored Error pr ocessing line 1 of /tmpfs/src/tf_docs_env/lib/python3.11/site-packages/tfdocs_sitecustomize.pth: Traceback (most recent call last): File "", line 195, in addpackage File "", line 1 import _tfdocs_sitecustomize\n ^ SyntaxError : unexpected character after line continuation character Remainder of file ignored WARNING: All log messages before absl::InitializeLog() is called are written to STDERR E0000 00:00:1793850186.812013 80889 cuda_platform.cc:57] failed call to cuInit: INTERNAL: CUDA error: Failed call to cuInit: CU DA_ERROR_NO_DEVICE: no CUDA-capable device is detected === Source Location Trace: === external/xla/ xla/stream_executor/cuda/cuda_status.cc:50 W0000 00:00:1793850198.183778 80889 dataset.cc:1000] I nput of GeneratorDatasetOp::Dataset will not be optimized because the dataset does not implement the AsGraphDefInternal() method needed to apply optimizations. Epoch 1/3 1/70 [...................... 2/70 [..............................] - ETA: 5s - loss: 4.615 4/70 [>....................... 5/70 [=>............................] - ETA: 4s - loss: 4.6295 - accuracy: 0.1688 7/70 [==>.......................... 8/70 [==>...........................] - ETA: 4s - loss: 4.6278 - accur acy: 0.1641 9/70 10/70 [===>..........................] - 11/70 [===>..........................] - ETA: 4s - loss: 4.6202 - accuracy: 0.1804 12/70 [==== 13/70 [====>.........................] - ETA: 14/70 [=====>........................] - ETA: 3s - loss: 4.6195 - accuracy: 0.179 7 15/70 [=====>... 16/70 [=====>........................] - ETA: 3s - 18/70 [======>....... 19/70 [=======>......................] - ETA: 3s - loss: 21/70 [========>.......... 22/70 [========>.....................] - ETA: 3s - loss: 4.61 24/70 [=========>.............. 25/70 [=========>....................] - ETA: 3s - loss: 4.6100 - accuracy: 0.1894 2 27/70 [==========>.................. 28/70 [===========>..................] - ETA: 3s - loss: 4.6081 - accur acy: 0.1931 29/70 30/70 [===========>..................] - 31/70 [============>.................] - ETA: 2s - loss: 4.6057 - accuracy: 0.2006 32/70 [==== 33/70 [=============>................] - ETA: 34/70 [=============>................] - ETA: 2s - loss: 4.6025 - accuracy: 0.204 5 35/70 [========= 36/70 [==============>...............] - ETA: 2s - 38/70 [============== 39/70 [===============>..............] - ETA: 2s - loss: 41/70 [================>.. 42/70 [=================>............] - ETA: 1s - loss: 4.59 44/70 [=================>...... 45/70 [==================>...........] - ETA: 1s - loss: 4.5920 - accuracy: 0.2198 4 47/70 [===================>......... 48/70 [===================>..........] - ETA: 1s - loss: 4.5896 - accur acy: 0.2253 49/70 50/70 [====================>.........] - 51/70 [====================>.........] - ETA: 1s - loss: 4.5873 - accuracy: 0.2313 52/70 [==== 53/70 [=====================>........] - ETA: 54/70 [======================>.......] - ETA: 1s - loss: 4.5851 - accuracy: 0.237 0 55/70 [========= 56/70 [=======================>......] - ETA: 0s - 58/70 [============== 59/70 [========================>.....] - ETA: 0s - loss: 61/70 [=================== 62/70 [=========================>....] - ETA: 0s - loss: 4.57 64/70 [======================== 65/70 [==========================>...] - ETA: 0s - loss: 4.5749 - accuracy: 0.2625 6 67/70 [===========================>. 68/70 [============================>.] - ETA: 0s - loss: 4.5723 - accur acy: 0.2693 69/70 70/70 [==============================] - 70/70 [==============================] - 6s 70ms/step - loss: 2.2850 - accur acy: 0.1368 Epoch 2/3 1/70 [..............................] - ETA: 4s - loss: 4.5224 - accuracy: 0 .4844 2/70 [..... 3/70 [>.............................] - ETA: 4 5/70 [=>........ 6/70 [=>............................] - ETA: 4s - l 8/70 [==>............ 9/70 [==>...........................] - ETA: 4s - loss: 11/70 [===>................ 12/70 [====>.........................] - ETA: 3s - loss: 4.496 14/70 [=====>................... 15/70 [=====>........................] - ETA: 3s - loss: 4.4937 - a ccuracy: 0.4688 16 17/70 [======>....................... 18/70 [======>.......................] - ETA: 3s - loss: 4.4926 - accura cy: 0.4714 19/70 [ 20/70 [=======>......................] - E 21/70 [========>.....................] - ETA: 3s - loss: 4.4919 - accuracy: 0 .4754 22/70 [===== 23/70 [========>.....................] - ETA: 3 25/70 [=========> 26/70 [==========>...................] - ETA: 2s - l 28/70 [===========>... 29/70 [===========>..................] - ETA: 2s - loss: 31/70 [============>....... 32/70 [============>.................] - ETA: 2s - loss: 4.483 34/70 [=============>........... 35/70 [==============>...............] - ETA: 2s - loss: 4.4829 - a ccuracy: 0.5147 36 37/70 [==============>............... 38/70 [===============>..............] - ETA: 2s - loss: 4.4814 - accura cy: 0.5185 39/70 [ 40/70 [================>.............] - E 41/70 [================>.............] - ETA: 1s - loss: 4.4787 - accuracy: 0 .5271 42/70 [===== 43/70 [=================>............] - ETA: 1 45/70 [========== 46/70 [==================>...........] - ETA: 1s - l 48/70 [=============== 49/70 [====================>.........] - ETA: 1s - loss: 51/70 [==================== 52/70 [=====================>........] - ETA: 1s - loss: 4.466 54/70 [======================>.. 55/70 [======================>.......] - ETA: 0s - loss: 4.4636 - a ccuracy: 0.5739 56 57/70 [=======================>...... 58/70 [=======================>......] - ETA: 0s - loss: 4.4606 - accura cy: 0.5846 59/70 [ 60/70 [========================>.....] - E 61/70 [=========================>....] - ETA: 0s - loss: 4.4571 - accuracy: 0 .5971 62/70 [===== 63/70 [==========================>...] - ETA: 0 65/70 [========== 66/70 [===========================>..] - ETA: 0s - l 68/70 [=============== 69/70 [============================>.] - ETA: 0s - loss: 70/70 [==================== ==========] - 5s 67ms/step - loss: 2.2243 - accuracy: 0.3095 Epoch 3/3 1/70 [..................... 2/70 [..............................] - ETA: 4s - loss: 4.3850 - accuracy: 0.8359 4/70 [>......................... 5/70 [=>............................] - ETA: 4s - loss: 4.3928 - ac curacy: 0.8094 6/ 7/70 [==>...........................] 8/70 [==>...........................] - ETA: 4s - loss: 4.3879 - accurac y: 0.7891 9/70 [= 10/70 [===>..........................] - ET 11/70 [===>..........................] - ETA: 3s - loss: 4.3815 - accuracy: 0. 8082 12/70 [====>. 13/70 [====>.........................] - ETA: 3s 15/70 [=====>..... 16/70 [=====>........................] - ETA: 3s - lo 18/70 [======>......... 19/70 [=======>......................] - ETA: 3s - loss: 4 21/70 [========>............ 22/70 [========>.....................] - ETA: 3s - loss: 4.3681 - accuracy: 0.8551 24/70 [=========>................ 25/70 [=========>....................] - ETA: 2s - loss: 4.3648 - ac curacy: 0.8556 26/ 27/70 [==========>...................] 28/70 [===========>..................] - ETA: 2s - loss: 4.3630 - accurac y: 0.8599 29/70 [= 30/70 [===========>..................] - ET 31/70 [============>.................] - ETA: 2s - loss: 4.3605 - accuracy: 0. 8669 32/70 [====== 33/70 [=============>................] - ETA: 2s 35/70 [=========== 36/70 [==============>...............] - ETA: 2s - lo 38/70 [===============> 39/70 [===============>..............] - ETA: 2s - loss: 4 41/70 [================>.... 42/70 [=================>............] - ETA: 1s - loss: 4.3485 - accuracy: 0.8932 44/70 [=================>........ 45/70 [==================>...........] - ETA: 1s - loss: 4.3454 - ac curacy: 0.8972 46/ 47/70 [===================>..........] 48/70 [===================>..........] - ETA: 1s - loss: 4.3439 - accurac y: 0.8975 49/70 [= 50/70 [====================>.........] - ET 51/70 [====================>.........] - ETA: 1s - loss: 4.3401 - accuracy: 0. 9050 52/70 [====== 53/70 [=====================>........] - ETA: 1s 55/70 [=========== 56/70 [=======================>......] - ETA: 0s - lo 58/70 [================ 59/70 [========================>.....] - ETA: 0s - loss: 4 61/70 [===================== 62/70 [=========================>....] - ETA: 0s - loss: 4.3287 - accuracy: 0.9302 64/70 [========================== 65/70 [==========================>...] - ETA: 0s - loss: 4.3245 - ac curacy: 0.9385 66/ 67/70 [===========================>..] 68/70 [============================>.] - ETA: 0s - loss: 4.3209 - accurac y: 0.9492 69/70 [= 70/70 [==============================] - ET 70/70 [==============================] - 5s 67ms/step - loss: 2.1596 - accurac y: 0.4751 ``` > [!NOTE] > **Note:** This may run slower than the test run at the beginning of this tutorial because running multiple workers on a single machine only adds overhead. The goal here is not to improve the training time but to give an example of multi-worker training. # Delete the `TF_CONFIG`, and kill any background tasks so they don't affect the next section. os.environ.pop('TF_CONFIG', None) %killbgscripts ``` All background processes were killed. ``` ## Multi-worker training in depth So far, you have learned how to perform a basic multi-worker setup. The rest of the tutorial goes over other factors, which may be useful or important for real use cases, in detail. ### Dataset sharding In multi-worker training, *dataset sharding* is needed to ensure convergence and performance. The example in the previous section relies on the default autosharding provided by the [`tf.distribute.Strategy`](https://www.tensorflow.org/api_docs/python/tf/distribute/Strategy) API. You can control the sharding by setting the [`tf.data.experimental.AutoShardPolicy`](https://www.tensorflow.org/api_docs/python/tf/data/experimental/AutoShardPolicy) of the [`tf.data.experimental.DistributeOptions`](https://www.tensorflow.org/api_docs/python/tf/data/experimental/DistributeOptions). To learn more about *auto-sharding* , refer to the [Distributed input guide](https://www.tensorflow.org/tutorials/distribute/input#sharding). Here is a quick example of how to turn the auto sharding off, so that each replica processes every example (*not recommended*): options = tf.data.Options() options.experimental_distribute.auto_shard_policy = tf.data.experimental.AutoShardPolicy.OFF global_batch_size = 64 multi_worker_dataset = mnist_setup.mnist_dataset(batch_size=64) dataset_no_auto_shard = multi_worker_dataset.with_options(options) ### Evaluation If you pass the `validation_data` into [`Model.fit`](https://www.tensorflow.org/api_docs/python/tf/keras/Model#fit) as well, it will alternate between training and evaluation for each epoch. The evaluation work is distributed across the same set of workers, and its results are aggregated and available to all workers. Similar to training, the validation dataset is automatically sharded at the file level. You need to set a global batch size in the validation dataset and set the `validation_steps`. A repeated dataset (by calling [`tf.data.Dataset.repeat`](https://www.tensorflow.org/api_docs/python/tf/data/Dataset#repeat)) is recommended for evaluation. Alternatively, you can also create another task that periodically reads checkpoints and runs the evaluation. This is what an Estimator does. But this is not a recommended way to perform evaluation and thus its details are omitted. ### Performance To tweak the performance of multi-worker training, you can try the following: - [`tf.distribute.MultiWorkerMirroredStrategy`](https://www.tensorflow.org/api_docs/python/tf/distribute/MultiWorkerMirroredStrategy) provides multiple [collective communication implementations](https://www.tensorflow.org/api_docs/python/tf/distribute/experimental/CommunicationImplementation): - `RING` implements ring-based collectives using gRPC as the cross-host communication layer. - `NCCL` uses the [NVIDIA Collective Communication Library](https://developer.nvidia.com/nccl) to implement collectives. - `AUTO` defers the choice to the runtime. The best choice of collective implementation depends upon the number of GPUs, the type of GPUs, and the network interconnects in the cluster. To override the automatic choice, specify the `communication_options` parameter of `MultiWorkerMirroredStrategy`'s constructor. For example: communication_options=tf.distribute.experimental.CommunicationOptions(implementation=tf.distribute.experimental.CommunicationImplementation.NCCL) - Cast the variables to `tf.float` if possible: - The official ResNet model includes [an example](https://github.com/tensorflow/models/blob/8367cf6dabe11adf7628541706b660821f397dce/official/resnet/resnet_model.py#L466) of how to do this. ### Fault tolerance In synchronous training, the cluster would fail if one of the workers fails and no failure-recovery mechanism exists. Using Keras with [`tf.distribute.Strategy`](https://www.tensorflow.org/api_docs/python/tf/distribute/Strategy) comes with the advantage of fault tolerance in cases where workers die or are otherwise unstable. You can do this by preserving the training state in the distributed file system of your choice, such that upon a restart of the instance that previously failed or preempted, the training state is recovered. When a worker becomes unavailable, other workers will fail (possibly after a timeout). In such cases, the unavailable worker needs to be restarted, as well as other workers that have failed. > [!NOTE] > **Note:** Previously, the `ModelCheckpoint` callback provided a mechanism to restore the training state upon a restart from a job failure for multi-worker training. The TensorFlow team is introducing a new [`BackupAndRestore`](https://www.tensorflow.org/tutorials/distribute/multi_worker_with_keras#scrollTo=kmH8uCUhfn4w) callback, which also adds the support to single-worker training for a consistent experience, and removed the fault tolerance functionality from existing `ModelCheckpoint` callback. From now on, applications that rely on this behavior should migrate to the new `BackupAndRestore` callback. #### The `ModelCheckpoint` callback `ModelCheckpoint` callback no longer provides fault tolerance functionality, please use [`BackupAndRestore`](https://www.tensorflow.org/tutorials/distribute/multi_worker_with_keras#scrollTo=kmH8uCUhfn4w) callback instead. The `ModelCheckpoint` callback can still be used to save checkpoints. But with this, if training was interrupted or successfully finished, in order to continue training from the checkpoint, the user is responsible to load the model manually. Optionally, users can choose to save and restore model/weights outside `ModelCheckpoint` callback. ### Model saving and loading To save your model using `model.save` or [`tf.saved_model.save`](https://www.tensorflow.org/api_docs/python/tf/saved_model/save), the saving destination needs to be different for each worker. - For non-chief workers, you will need to save the model to a temporary directory. - For the chief, you will need to save to the provided model directory. The temporary directories on the worker need to be unique to prevent errors resulting from multiple workers trying to write to the same location. The model saved in all the directories is identical, and typically only the model saved by the chief should be referenced for restoring or serving. You should have some cleanup logic that deletes the temporary directories created by the workers once your training has completed. The reason for saving on the chief and workers at the same time is because you might be aggregating variables during checkpointing, which requires both the chief and workers to participate in the allreduce communication protocol. On the other hand, letting chief and workers save to the same model directory will result in errors due to contention. Using the `MultiWorkerMirroredStrategy`, the program is run on every worker, and in order to know whether the current worker is the chief, it takes advantage of the cluster resolver object that has attributes `task_type` and `task_id`: - `task_type` tells you what the current job is (for example, `'worker'`). - `task_id` tells you the identifier of the worker. - The worker with `task_id == 0` is designated as the chief worker. In the code snippet below, the `write_filepath` function provides the file path to write, which depends on the worker's `task_id`: - For the chief worker (with `task_id == 0`), it writes to the original file path. - For other workers, it creates a temporary directory---`temp_dir`---with the `task_id` in the directory path to write in: model_path = '/tmp/keras-model' def _is_chief(task_type, task_id): # Note: there are two possible `TF_CONFIG` configurations. # 1) In addition to `worker` tasks, a `chief` task type is use; # in this case, this function should be modified to # `return task_type == 'chief'`. # 2) Only `worker` task type is used; in this case, worker 0 is # regarded as the chief. The implementation demonstrated here # is for this case. # For the purpose of this Colab section, the `task_type` is `None` case # is added because it is effectively run with only a single worker. return (task_type == 'worker' and task_id == 0) or task_type is None def _get_temp_dir(dirpath, task_id): base_dirpath = 'workertemp_' + str(task_id) temp_dir = os.path.join(dirpath, base_dirpath) tf.io.gfile.makedirs(temp_dir) return temp_dir def write_filepath(filepath, task_type, task_id): dirpath = os.path.dirname(filepath) base = os.path.basename(filepath) if not _is_chief(task_type, task_id): dirpath = _get_temp_dir(dirpath, task_id) return os.path.join(dirpath, base) task_type, task_id = (strategy.cluster_resolver.task_type, strategy.cluster_resolver.task_id) write_model_path = write_filepath(model_path, task_type, task_id) With that, you're now ready to save: > [!WARNING] > **Deprecated:** For Keras objects, it's recommended to use the new high-level `.keras` format and [`tf.keras.Model.export`](https://www.tensorflow.org/api_docs/python/tf/keras/Model#export), as demonstrated in the guide [here](https://www.tensorflow.org/guide/keras/save_and_serialize). The low-level SavedModel format continues to be supported for existing code. multi_worker_model.save(write_model_path) ``` INFO:tensorflow:Assets written to: /tmp/keras-model/assets INFO:tensorflow:Assets written to: /tmp/keras-model/assets ``` As described above, later on the model should only be loaded from the file path the chief worker saved to. Therefore, remove the temporary ones the non-chief workers have saved: if not _is_chief(task_type, task_id): tf.io.gfile.rmtree(os.path.dirname(write_model_path)) Now, when it's time to load, use the convenient [`tf.keras.models.load_model`](https://www.tensorflow.org/api_docs/python/tf/keras/models/load_model) API, and continue with further work. Here, assume only using single worker to load and continue training, in which case you do not call [`tf.keras.models.load_model`](https://www.tensorflow.org/api_docs/python/tf/keras/models/load_model) within another `strategy.scope()` (note that `strategy = tf.distribute.MultiWorkerMirroredStrategy()`, as defined earlier): loaded_model = tf.keras.models.load_model(model_path) # Now that the model is restored, and can continue with the training. loaded_model.fit(single_worker_dataset, epochs=2, steps_per_epoch=20) ``` Epoch 1/2 20/20 [==============================] - 1s 7ms/step - loss: 2.3102 - accuracy: 7.8125e-04 Epoch 2/2 20/20 [==============================] - 0s 7ms/step - loss: 2.3041 - accuracy: 0.0016 ``` ### Checkpoint saving and restoring On the other hand, checkpointing allows you to save your model's weights and restore them without having to save the whole model. Here, you'll create one [`tf.train.Checkpoint`](https://www.tensorflow.org/api_docs/python/tf/train/Checkpoint) that tracks the model, which is managed by the [`tf.train.CheckpointManager`](https://www.tensorflow.org/api_docs/python/tf/train/CheckpointManager), so that only the latest checkpoint is preserved: checkpoint_dir = '/tmp/ckpt' checkpoint = tf.train.Checkpoint(model=multi_worker_model) write_checkpoint_dir = write_filepath(checkpoint_dir, task_type, task_id) checkpoint_manager = tf.train.CheckpointManager( checkpoint, directory=write_checkpoint_dir, max_to_keep=1) Once the `CheckpointManager` is set up, you're now ready to save and remove the checkpoints the non-chief workers had saved: checkpoint_manager.save() if not _is_chief(task_type, task_id): tf.io.gfile.rmtree(write_checkpoint_dir) Now, when you need to restore the model, you can find the latest checkpoint saved using the convenient [`tf.train.latest_checkpoint`](https://www.tensorflow.org/api_docs/python/tf/train/latest_checkpoint) function. After restoring the checkpoint, you can continue with training. latest_checkpoint = tf.train.latest_checkpoint(checkpoint_dir) checkpoint.restore(latest_checkpoint) multi_worker_model.fit(multi_worker_dataset, epochs=2, steps_per_epoch=20) ``` Epoch 1/2 W0000 00:00:1793850216.634229 80694 dataset.cc:1000] Input of GeneratorDatasetOp::Dataset will not be optimized because the dataset does not implement the AsGraphDefInternal() method needed to apply optimizations. 20/20 [==============================] - 1s 9ms/step - loss: 2.3108 - accuracy: 0.0000e+00 Epoch 2/2 20/20 [==============================] - 0s 9ms/step - loss: 2.3068 - accuracy: 7.8125e-04 ``` #### The `BackupAndRestore` callback The [`tf.keras.callbacks.BackupAndRestore`](https://www.tensorflow.org/api_docs/python/tf/keras/callbacks/BackupAndRestore) callback provides the fault tolerance functionality by backing up the model and current training state in a temporary checkpoint file under `backup_dir` argument to `BackupAndRestore`. > [!NOTE] > **Note:** In Tensorflow 2.9, the current model and the training state is backed up at epoch boundaries. In the `tf-nightly` version and from TensorFlow 2.10, the `BackupAndRestore` callback can back up the model and the training state at epoch or step boundaries. `BackupAndRestore` accepts an optional `save_freq` argument. `save_freq` accepts either `'epoch'` or an `int` value. If `save_freq` is set to `'epoch'` the model is backed up after every epoch. If `save_freq` is set to an integer value greater than `0`, the model is backed up after every `save_freq` number of batches. Once the jobs get interrupted and restarted, the `BackupAndRestore` callback restores the last checkpoint, and you can continue training from the beginning of the epoch and step at which the training state was last saved. To use it, provide an instance of [`tf.keras.callbacks.BackupAndRestore`](https://www.tensorflow.org/api_docs/python/tf/keras/callbacks/BackupAndRestore) at the [`Model.fit`](https://www.tensorflow.org/api_docs/python/tf/keras/Model#fit) call. With `MultiWorkerMirroredStrategy`, if a worker gets interrupted, the whole cluster will pause until the interrupted worker is restarted. Other workers will also restart, and the interrupted worker will rejoin the cluster. Then, every worker will read the checkpoint file that was previously saved and pick up its former state, thereby allowing the cluster to get back in sync. Then, the training will continue. The distributed dataset iterator state will be re-initialized and not restored. The `BackupAndRestore` callback uses the `CheckpointManager` to save and restore the training state, which generates a file called checkpoint that tracks existing checkpoints together with the latest one. For this reason, `backup_dir` should not be re-used to store other checkpoints in order to avoid name collision. Currently, the `BackupAndRestore` callback supports single-worker training with no strategy---`MirroredStrategy`---and multi-worker training with `MultiWorkerMirroredStrategy`. Below are two examples for both multi-worker training and single-worker training: # Multi-worker training with `MultiWorkerMirroredStrategy` # and the `BackupAndRestore` callback. The training state # is backed up at epoch boundaries by default. callbacks = [tf.keras.callbacks.BackupAndRestore(backup_dir='/tmp/backup')] with strategy.scope(): multi_worker_model = mnist_setup.build_and_compile_cnn_model() multi_worker_model.fit(multi_worker_dataset, epochs=3, steps_per_epoch=70, callbacks=callbacks) ``` Epoch 1/3 70/70 [==============================] - 1s 10ms/step - loss: 2.2885 - accuracy: 0.1623 Epoch 2/3 70/70 [==============================] - 1s 10ms/step - loss: 2.2335 - accuracy: 0.3225 Epoch 3/3 70/70 [==============================] - 1s 10ms/step - loss: 2.1680 - accuracy: 0.5158 ``` If the `save_freq` argument in the `BackupAndRestore` callback is set to `'epoch'`, the model is backed up after every epoch. # The training state is backed up at epoch boundaries because `save_freq` is # set to `epoch`. callbacks = [tf.keras.callbacks.BackupAndRestore(backup_dir='/tmp/backup')] with strategy.scope(): multi_worker_model = mnist_setup.build_and_compile_cnn_model() multi_worker_model.fit(multi_worker_dataset, epochs=3, steps_per_epoch=70, callbacks=callbacks) ``` Epoch 1/3 70/70 [==============================] - 1s 11ms/step - loss: 2.2764 - accuracy: 0.1737 Epoch 2/3 70/70 [==============================] - 1s 11ms/step - loss: 2.2126 - accuracy: 0.3529 Epoch 3/3 70/70 [==============================] - 1s 11ms/step - loss: 2.1295 - accuracy: 0.5100 ``` > [!NOTE] > **Note:** The next code block uses features that are only available in `tf-nightly` until Tensorflow 2.10 is released. If the `save_freq` argument in the `BackupAndRestore` callback is set to an integer value greater than `0`, the model is backed up after every `save_freq` number of batches. # The training state is backed up at every 30 steps because `save_freq` is set # to an integer value of `30`. callbacks = [tf.keras.callbacks.BackupAndRestore(backup_dir='/tmp/backup', save_freq=30)] with strategy.scope(): multi_worker_model = mnist_setup.build_and_compile_cnn_model() multi_worker_model.fit(multi_worker_dataset, epochs=3, steps_per_epoch=70, callbacks=callbacks) ``` Epoch 1/3 70/70 [==============================] - 1s 12ms/step - loss: 2.2689 - accuracy: 0.2020 Epoch 2/3 70/70 [==============================] - 1s 12ms/step - loss: 2.2051 - accuracy: 0.4342 Epoch 3/3 70/70 [==============================] - 1s 13ms/step - loss: 2.1248 - accuracy: 0.5641 ``` If you inspect the directory of `backup_dir` you specified in `BackupAndRestore`, you may notice some temporarily generated checkpoint files. Those files are needed for recovering the previously lost instances, and they will be removed by the library at the end of [`Model.fit`](https://www.tensorflow.org/api_docs/python/tf/keras/Model#fit) upon successful exiting of your training. > [!NOTE] > **Note:** Currently the `BackupAndRestore` callback only supports eager mode. In graph mode, consider using [`Model.save`](https://www.tensorflow.org/api_docs/python/tf/keras/Model#save)/[`tf.saved_model.save`](https://www.tensorflow.org/api_docs/python/tf/saved_model/save) and [`tf.keras.models.load_model`](https://www.tensorflow.org/api_docs/python/tf/keras/models/load_model) for saving and restoring models, respectively, as described in the *Model saving and loading* section above, and by providing `initial_epoch` in [`Model.fit`](https://www.tensorflow.org/api_docs/python/tf/keras/Model#fit) during training. ## Additional resources 1. The [Distributed training in TensorFlow](https://www.tensorflow.org/guide/distributed_training) guide provides an overview of the available distribution strategies. 2. The [Custom training loop with Keras and MultiWorkerMirroredStrategy](https://www.tensorflow.org/tutorials/distribute/multi_worker_with_ctl) tutorial shows how to use the `MultiWorkerMirroredStrategy` with Keras and a custom training loop. 3. Check out the [official models](https://github.com/tensorflow/models/tree/master/official), many of which can be configured to run multiple distribution strategies. 4. The [Better performance with tf.function](https://www.tensorflow.org/guide/function) guide provides information about other strategies and tools, such as the [TensorFlow Profiler](https://www.tensorflow.org/guide/profiler) you can use to optimize the performance of your TensorFlow models.