0%

Mixed-Precision Training of Deep Neural Networks

由于FP16的精度损失问题,如果我们在神经网络的训练过程中直接将网络参数以FP16的形式进行计算,可能会出现数值不稳定的情况而导致模型性能下降。 对此,百度和NVIDIA的研究院在Mixed Precision Training这一论文中提出混合精度训练的方法,在充分利用FP16加速运算的优点的同时保证了模型的精度。下面介绍其论文中的主要要点。

FP32 master copy of weigths

FP32 master copy即维护一份网络中FP16精度参数的FP32精度的拷贝。 计算过程如下图所示,在前向传播过程中,使用由master copy类型转换得到的FP16精度参数进行运算; 而在反向传播计算完梯度后,将梯度作用到master copy上以在FP32精度上进行参数更新。

FP32 master copyFP32 master copy

这么做的主要原因有两个。

第一是由于若直接在FP16精度下进行参数更新,存在梯度过小而导致更新值为0的情况。 下图是某次模型训练过程汇总网络参数的梯度的分布直方图,有5%的梯度值是分布在小于2-24的区间内的。 若优化器直接将这部分梯度乘上学习率作用到FP16精度参数的更新上,那么更新值将为0。 这将影响到模型的准确率。 但是如果是更新值是作用到FP32精度的master copy上时,则不会出现更新值下溢为0的情况。

第二是如果参数相较于其更新值过大的话,也可能会由于浮点数加法机制的缘故而导致更新值为0。 在浮点数加法过程中,需要将两数对齐进行运算。 如果参数的大小是其更新值的2048倍或者更大的话,那么更新值的小数位需要右移至少11位才能与前者对齐,这超出了FP16精度的表示范围。 在FP32精度下则一般不会出现这种问题。

  1. gradient histogramgradient histogram

  2. loss scaling

    loss scaling即将loss值放大,以保证反向传播过程当中梯度落在FP16精度能表示的范围之间。

    下图是Multibox SSD网络在训练过程过程中激活单元梯度的分布情况。其中有67%的梯度落在了小于2-24的范围内,在FP16精度下无法表示。 如果不对梯度进行放大,在FP16精度下对该网络进行训练将导致发散。 对激活梯度进行放大,再对参数梯度缩小相应倍数,即可解决该问题。

    activation gradientsactivation gradients

    根据梯度计算的链式法则,对梯度放大的最简单的方法就是放大loss。 放大因子的大小选择没有固定的标准,对于上述Multibox SSD网络,作者尝试了8-32K的放大因子,均训练成功了。 只要保证放大后的梯度不超过FP16精度的表示上限(65504),选择较大的放大因子并无副作用。[1]

  1. Nvidia Reources: https://github.com/NvidiaResources/nvidia_mixed_precision_training
  2. Nvidia: Mixed-Precision Training of Deep Neural Networks: https://devblogs.nvidia.com/mixed-precision-training-deep-neural-networks/

3. Method

2.1 直接在python2的pytorch训练框架中修改可实现:

https://github.com/suvojit-0x55aa/mixed-precision-pytorch

2.2 使用Nvidia apex库

库地址:https://github.com/NVIDIA/apex

Nvidia 官方apex使用教程:http://on-demand.gputechconf.com/gtc-cn/2018/pdf/CH8302.pdf

Reference

  1. 混合精度训练