Background¶

The computational resources required for training deep neural networks have been increasing of late because of complexity of the architectures and size of models. Mixed precision training allows us to reduces the resources required by using lower precision arithmetic. In this approach you can train using 16 bit floating points (half precision) while using 32 bit floating points (single precision) for output buffers of float16 computation. This combination of single and half precision gives rise to the name mixed precision. It allows us to achieve the same accuracy as training with single precision, while decreasing the required memory and training or inference time.

The float16 data type is a 16 bit floating point representation according to the IEEE 754 standard. It has a dynamic range where the precision can go from 0.0000000596046 (highest, for values closest to 0) to 32 (lowest, for values in the range 32768-65536). Despite the inherent reduced precision when compared to single precision float (float32), using float16 has many advantages. The most obvious advantages are that you can reduce the size of the model by half allowing the training of larger models and using larger batch sizes. The reduced memory footprint also helps in reducing the pressure on memory bandwidth and lowering communication costs. On hardware with specialized support for float16 computation you can also greatly improve the speed of training and inference. The Volta range of Graphics Processing Units (GPUs) from Nvidia have Tensor Cores which perform efficient float16 computation. A tensor core allows accumulation of half precision products into single or half precision outputs. For the rest of this tutorial we assume that we are working with Nvidia’s Tensor Cores on a Volta GPU.

Prerequisites¶

• Volta range of Nvidia GPUs
• Cuda 9 or higher
• CUDNN v7 or higher

This tutorial also assumes that you understand how to train a network with float32. Please refer to other tutorials here to get started with MXNet and/or Gluon. This tutorial focuses on the changes needed to switch from float32 to mixed precision and tips on achieving the best performance with mixed precision.

Using the Gluon API¶

net = net.cast('float16')

data = data.astype('float16', copy=False)

optimizer = mx.optimizer.create('sgd', multi_precision=True, lr=0.01)

python image_classification.py --model resnet50_v1 --dataset caltech101 --gpus 0 --num-worker 30 --dtype float16

pretrained_net = models.get_model(name='resnet50_v2', ctx=ctx, pretrained=True, classes=1000)
pretrained_net.cast('float16')

net = models.get_model(name='resnet50_v2', ctx=ctx, pretrained=False, classes=101)
net.collect_params().initialize(mx.init.Xavier(magnitude=2.24), ctx=ctx)
net.features = pretrained_net.features
net.cast(dtype)

Using the Symbolic API¶

Training a network in float16 with the Symbolic API involves the following steps.

1. Add a layer at the beginning of the network, to cast the data to float16. This will ensure that all the following layers compute in float16.
2. It is advisable to cast the output of the layers before softmax to float32, so that the softmax computation is done in float32. This is because softmax involves large reductions and it helps to keep that in float32 for more precise answer.
3. It is advisable to use the multi-precision mode of the optimizer for more precise weight updates. This is discussed in some detail below. Here’s how you would enable this mode when creating an optimizer.

optimizer = mx.optimizer.create('sgd', multi_precision=True, lr=0.01)

There are a few examples of building such networks which can handle float16 input in examples/image-classification/symbols/. Specifically you could look at the resnet example.

An illustration of the relevant section of the code is below.

data = mx.sym.Variable(name="data")
if dtype == 'float16':
    data = mx.sym.Cast(data=data, dtype=np.float16)

// the rest of the network 
net_out = net(data)

if dtype == 'float16':
    net_out = mx.sym.Cast(data=net_out, dtype=np.float32)
output = mx.sym.SoftmaxOutput(data=net_out, name='softmax')

We have an example script which show how to train imagenet with resnet50 using float16 here

Here’s how you can use the above script to train Resnet50 v1 model with synthetic data using float16, so you can try it out even if you don’t have the Imagenet dataset handy.

python train_imagenet.py --network resnet-v1 --num-layers 50 --benchmark 1 --gpus 0 --batch-size 256 --dtype float16

There’s a similar example for fine tuning here. The following command shows how to use that script to fine tune a Resnet50 model trained on Imagenet for the Caltech 256 dataset using float16.

python fine-tune.py --network resnet --num-layers 50 --pretrained-model imagenet1k-resnet-50 --data-train ~/data/caltech-256/caltech256-train.rec --data-val ~/data/caltech-256/caltech256-val.rec --num-examples 15420 --num-classes 256 --gpus 0 --batch-size 64 --dtype float16

Example training results¶

Let us consider training a Resnet50 v1 model on the Imagenet 2012 dataset. For this model, the GPU memory usage is close to the capacity of V100 GPU with a batch size of 128 when using float32. Using float16 allows the use of 256 batch size. Shared below are results using 8 V100 GPUs on a AWS p3.16x large instance. Let us compare the three scenarios that arise here: float32 with 1024 batch size, float16 with 1024 batch size and float16 with 2048 batch size. These jobs trained for 90 epochs using a learning rate of 0.4 for 1024 batch size and 0.8 for 2048 batch size. This learning rate was decayed by a factor of 0.1 at the 30th, 60th and 80th epochs. The only changes made for the float16 jobs when compared to the float32 job were that the network and data were cast to float16, and the multi-precision mode was used for optimizer. The final accuracy at 90th epoch and the time to train are tabulated below for these three scenarios. The top-1 validation errors at the end of each epoch are also plotted below.

Batch size | Data type | Top 1 Validation accuracy | Time to train | Speedup | — | — | — | — | — | 1024 | float32 | 76.18% | 11.8 hrs | 1 | 1024 | float16 | 76.34% | 7.3 hrs | 1.62x | 2048 | float16 | 76.29% | 6.5 hrs | 1.82x |

The differences in accuracies above are within normal random variation, and there is no reason to expect float16 to have better accuracy than float32 in general. As the plot indicates training behaves similarly for these cases, even though we didn’t have to change any other hyperparameters. We can also see from the table that using float16 helps train faster through faster computation with float16 as well as allowing the use of larger batch sizes.

Things to keep in mind¶

loss = gluon.loss.SoftmaxCrossEntropyLoss(weight=128)
optimizer = mx.optimizer.create('sgd', multi_precision=True, rescale_grad=1.0/128)

mxnet.sym.SoftmaxOutput(other_args, grad_scale=128.0)
optimizer = mx.optimizer.create('sgd', multi_precision=True, rescale_grad=1.0/128)

Video Tutorial¶

We also have a video tutorial for using Mixed Precision with MXNet. You can check that out here

References¶