Train a Neural Network on multi-GPU with TensorFlow | by Jordi TORRES.AI | Towards Data Science

[This post will be used in the master course Supercomputers Architecture at UPC Barcelona Tech with the support of the BSC]
In the previous post we introduced the basic concepts in distributed and parallel training Deep Learning. We also presented the frameworks and infrastructure that we will use to scale the learning process of different neural networks for classifying the popular CIFAR10 dataset.
Now, we are going to explore how we can scale the training on Multiple GPUs in one Server with TensorFlow using tf.distributed.MirroredStrategy().
tf.distribute.Strategy is a TensorFlow API to distribute training across multiple GPU or TPUs with minimal code changes (from the sequential version presented in the previous post). This API can be used with a high-level API like Keras, and can also be used to distribute custom training loops.
tf.distribute.Strategy intends to cover a number of distribution strategies use cases along different axes. The official web page of this feature, presents all the currently supported combinations, however, in this article we will focus our attention in tf.distribute.MirroredStrategy.
We will use in this course tf.distribute.MirroredStrategy , that supports synchronous distributed training on multiple GPUs on one server. It creates one replica per GPU device. Each variable in the model is mirrored across all the replicas. These variables are kept in sync with each other by applying identical updates.
Efficient all-reduce algorithms presented in the previous post are used to communicate the variable updates across the GPUs. By default, it uses NVIDIA NCCL as the all-reduce implementation.
Here is the simplest way of creating MirroredStrategy:
mirrored_strategy = tf.distribute.MirroredStrategy()
This will create a MirroredStrategy instance that will use all the GPUs visible to TensorFlow and use NCCL as the cross-device communication. It is possible to see the list of available GPU devices doing the following:
devices = tf.config.experimental.list_physical_devices(“GPU”)
It is also possible to use a subset of the available GPUs in the system by doing the following:
mirrored_strategy = tf.distribute.MirroredStrategy(devices=["/gpu:0", "/gpu:1"])
To build the model and compile it inside the MirroredStrategyscope you can do it in the following way:
with mirrored_strategy.scope():     model = tf.keras.applications.resnet_v2.ResNet50V2(             include_top=True, weights=None,              input_shape=(128, 128, 3), classes=10)       opt = tf.keras.optimizers.SGD(learning_rate)      model.compile(loss=’sparse_categorical_crossentropy’,                    optimizer=opt, metrics=[‘accuracy’])
This allows us to create distributed variables instead of regular variables: each variable is mirrored across all the replicas and are kept in sync with each other by applying identical updates.
Training can be done as usual outside the strategy scope with:
dataset = load_data(batch_size)model.fit(dataset, epochs=5, verbose=2)
In this post, we will consider using the epoch time as a measure of the computation time for training a DNN. This approximated measure in seconds provided by Keras during the .fit is enough for the purpose of this academic post. In our case, sometimes we will discard the first time epoch as it includes creating and initializing structures. Obviously, for certain types of studies, it is necessary to go into more detail, differentiating the loading data, feeds forward time, loss function time, back progatation time, etc., but it falls outside the scope of this case study we propose here.
When training, it is required to allocate memory to store samples for training the model and the model itself. We have to keep in mind this in order to avoid an out of memory error.
Remember that the batch_size is the number of samples that the model will see at each training step, and in general, we want to have this number as biggest as possible (powers of 2). We can calculate it by try and error approach testing different values until an error related to the memory appears:
python ResNet50.py -- epoch 1 -- batch_size 16python ResNet50.py -- epoch 1 -- batch_size 32python ResNet50.py -- epoch 1 -- batch_size 64
When using MirroredStrategy with multiple GPUs, the batch size indicated is divided by the number of replicas. Therefore the batch_size that we should specify to TensorFlow is equal to the maximum value for one GPU multiplied by the number of GPUs we are using. This is, in our example, use these flags in the python program:
python ResNet50.py -- epochs 5 -- batch_size 256  -- n_gpus 1python ResNet50.py -- epochs 5 -- batch_size 512  -- n_gpus 2python ResNet50.py -- epochs 5 -- batch_size 1024 -- n_gpus 4
Accordingly, with the batch_size, if we are using MirroredStrategy with multiple GPUs, we want the learning_rate to be:
learning_rate = learning_rate_base*number_of_gpusopt = tf.keras.optimizers.SGD(learning_rate)
Because of a larger batch_size, we also want to take bigger steps in the direction of the minimum to preserve the number of epochs to converge.
We will take the example presented in the previous post, ResNet50 with CIFAR10, and train it on the CTE-POWER machine.
Following the steps presented in the above section on how to apply MirroredStrategy, below we present the resulting parallel code:
import tensorflow as tf from tensorflow.keras import layers from tensorflow.keras import modelsimport numpy as npimport argparseimport timeimport syssys.path.append(‘/gpfs/projects/nct00/nct00002/cifar-utils’)from cifar import load_cifarparser = argparse.ArgumentParser()parser.add_argument(‘ -- epochs’, type=int, default=5)parser.add_argument(‘ -- batch_size’, type=int, default=2048)parser.add_argument(‘ -- n_gpus’, type=int, default=1)args = parser.parse_args()batch_size = args.batch_sizeepochs = args.epochsn_gpus = args.n_gpustrain_ds, test_ds = load_cifar(batch_size)device_type = ‘GPU’devices = tf.config.experimental.list_physical_devices(          device_type)devices_names = [d.name.split(“e:”)[1] for d in devices]strategy = tf.distribute.MirroredStrategy(           devices=devices_names[:n_gpus])with strategy.scope():     model = tf.keras.applications.resnet_v2.ResNet50V2(             include_top=True, weights=None,             input_shape=(128, 128, 3), classes=10)     opt = tf.keras.optimizers.SGD(0.01*n_gpus)     model.compile(loss=’sparse_categorical_crossentropy’,                    optimizer=opt, metrics=[‘accuracy’])model.fit(train_ds, epochs=epochs, verbose=2)
The SLURM script to send the jobs to execution is:
#!/bin/bash#SBATCH --job-name=”ResNet50"#SBATCH --D .#SBATCH --output=ResNet50_%j.output#SBATCH --error=ResNet50_%j.err#SBATCH --nodes=1#SBATCH --ntasks=1#SBATCH --cpus-per-task=160#SBATCH --gres=gpu:4#SBATCH --time=00:20:00module purge; module load gcc/8.3.0 cuda/10.2 cudnn/7.6.4 nccl/2.4.8 tensorrt/6.0.1 openmpi/4.0.1 atlas/3.10.3 scalapack/2.0.2 fftw/3.3.8 szip/2.1.1 ffmpeg/4.2.1 opencv/4.1.1 python/3.7.4_MLpython ResNet50.py -- epochs 5 -- batch_size 256 -- n_gpus 1python ResNet50.py -- epochs 5 -- batch_size 512 -- n_gpus 2python ResNet50.py -- epochs 5 -- batch_size 1024 -- n_gpus 4
Once I run the previous script, in the file that has stored the standard output we find the following execution times:
python ResNet50.py --epochs 5 --batch_size 256 --n_gpus 1 Epoch 1/5 196/196 - 49s - loss: 2.0408 - accuracy: 0.2506Epoch 2/5196/196 - 45s - loss: 1.7626 - accuracy: 0.3536Epoch 3/5196/196 - 45s - loss: 1.5863 - accuracy: 0.4164Epoch 4/5196/196 - 45s - loss: 1.4550 - accuracy: 0.4668Epoch 5/5196/196 - 45s - loss: 1.3539 - accuracy: 0.5070python ResNet50.py --epochs 5 --batch_size 512 --n_gpus 2Epoch 1/598/98 - 26s - loss: 2.0314 - accuracy: 0.2498Epoch 2/598/98 - 24s - loss: 1.7187 - accuracy: 0.3641Epoch 3/598/98 - 24s - loss: 1.5731 - accuracy: 0.4207Epoch 4/598/98 - 23s - loss: 1.4543 - accuracy: 0.4686Epoch 5/598/98 - 23s - loss: 1.3609 - accuracy: 0.5049 python ResNet50.py --epochs 5 --batch_size 1024 --n_gpus 4 Epoch 1/549/49 - 14s - loss: 2.0557 - accuracy: 0.2409Epoch 2/549/49 - 12s - loss: 1.7348 - accuracy: 0.3577Epoch 3/549/49 - 12s - loss: 1.5696 - accuracy: 0.4180Epoch 4/549/49 - 12s - loss: 1.4609 - accuracy: 0.4625Epoch 5/549/49 - 11s - loss: 1.3689 - accuracy: 0.5010 
It is important to note that we center our interest in the computational speed of the process rather than the model’s accuracy. For this reason, we will only execute a few number of epochs during the training, due as we will see training time per epoch are constant, and with 5 epochs we achieve the same accuracy in all cases. That means that only 5 epochs serve us to compare the 3 options, the purpose of these tests.
It is evident that if we want to do a good study, we should carry out several tests and then take the averages of the times obtained in each one of them. But given the academic purpose of this post (and the cost of resources, which we have to save!), it is enough only one execution.
In summary, the time required for doing one epoch is:
We can translate this to images per second (since we know there are 50,000 images), which gives us the throughput:
Finally, as I said, the speedup is the most relevant metric:
It’s an almost linear speedup. We refer to linear speedup linear when the workload is equally divided between the number of GPUs.
If you want to learn more and consolidate the knowledge acquired, now it’s your turn to get your hands dirty to reproduce the above results for the ResNET152V2, and then compare them with those presented here. Remember that the first step is to find the best batch_size
If you plot the results, you should find results equivalent to those shown below:
A couple of relevant things are observed. The first is that the time to run a ResNet152 epoch is much longer, and therefore the throughput in images per second is much lower than on the ResNet50 network. Why is this happening?
It can be seen that the speedup for the ResNet152 is no longer linear; could you try to give a possible answer as to why this is happening?
In this post, we have presented how we could paralyze a single Deep Neural Network training over many GPUs in one server using tf.distributed.MirroredStrategy() TensorFlow API.
As we can see in the results, the accuracy achieved by our model is more or less constant independently of the number of GPUs. In summary, using a distributed strategy, our model can learn the same but faster, thanks to parallelization!
In the next post, we will show how you can scale using multiple servers using the Horovod package. See you in the next post.
Thanks to parallelization, DL can learn the same but faster!
Update 15/12/2020: TensorFlow has released its new update 2.4 with features with increased support for parallel training (link)
Artificial Intelligence is a Supercomputing problemUsing Supercomputers for Deep Learning TrainingScalable Deep Learning on Parallel and Distributed InfrastructuresTrain a Neural Network on multi-GPU with TensorFlowDistributed Deep Learning with Horovod
Artificial Intelligence is a Supercomputing problem
Using Supercomputers for Deep Learning Training
Scalable Deep Learning on Parallel and Distributed Infrastructures
Train a Neural Network on multi-GPU with TensorFlow
Distributed Deep Learning with Horovod
Acknowledgment: Many thanks to Juan Luis Domínguez and Oriol Aranda from BSC-CNS, who wrote the first version of the codes that appear in this series of posts, and to Carlos Tripiana for the essential support with the deployment of the software stack of POWER-CTE Supercomputer. Also many thanks to Alvaro Jover Alvarez, Miquel Escobar Castells and Raul Garcia Fuentes for their contributions to the proofreading of this document.