混合精度训练深度神经网络

背景

随着深度神经网络的快速发展，网络结构和数据集呈现复杂化，巨大化；如何快速的训练一个深度神经网络面临着严峻的挑战。NVIDIA最新发布的带有Tensor Core的GPU，如V100, P4, P40, P100等卡可以支持单精度(FP32)和半精度(FP16)的混合训练，混合训练中以半精度为主，单精度为辅，可以在保持网络性能的同时，大大提高网络训练速度和推理速度和显存。

机器和软件环境

• V系列或者P系列等含有Tensor Core的GPU
• CUDA 9及以上
• cuDNN v7及以上

IEEE 754 FP16的表示范围

FP16可以表示正规形式或者非正规形式，非正规形式能够表示绝对值更接近0的数字。

• 正规形式
  在IEEE754的标准，16位的浮点数，最高位表示符号域，接着的5位表示指数域，最后的10位表示尾数域。

  • 指数域的取值范围为[-2^{e-1} + 1, 2^{e-1} + 1], 即[-14, 15]。
  • 尾数域的取值范围，在尾数表示中，会隐藏小数点前的1，即尾数最小值为1, 最大值为2-2^10
  • 那么整个FP16能够表示数字的绝对值范围是[2^{-14} * 1, (2-2^10) * 2^15]， 即[6.10e-5, 65504]

• 非正规形式
  在上述正规形式中，尾数隐藏了一个1，即尾数的最小值为1，那么也可以让尾数表示一个小于1的小数，并且隐藏小数点前的0，并且规定指数域全为0

  • 最小的数字的绝对值为2^{-24} ~= 5.96e-8
  • 最大的数字的绝对值为(1-2^{-10})*2^{-14} ~= 6.10e-5

因此，FP16一共能表示2^40个数字，绝对值范围为[5.96e-8, 65504]

实战

当前的大多数深度学习框架已经继承了NVIDIA的半精度训练，能够做到只修改少量代码就能够支持混合精度训练，例如：PyTorch, TensorFlow, MXNet, Caffe2, CNTK, NVCaffe, Theano。
以MXNet为例说明如何将代码修改到支持混合精度训练和推理：

• 训练

  from mxnet.contrib import amp
  amp.init()
  amp.init_trainer(trainer)

  上述代码可以用到Symbol和Gluon两种模式的训练，不过有以下限制要注意：

  1. 动态scale只调整支持Gluon trainer而且update_on_kvstore = False
  2. 使用SoftmaxOutput, LinearRegressionOutput, LogisticRegressionOutput,MAERegressionOutput时不能有多个Gluon trainer
• 推理
  如果我们已经有一个训练好的模型，想要在部署的时候，使用半精度或者混合精度推理，加快推理速度，节省内存，可以通过amp.convert_model对Symbol模型进行转换，amp.convert_hybrid_block对Gluon Hybrid模型进行转换，这两个API支持将一部分网络层转换为FP16，另一部分仍旧保持FP32。

AMP内部原理

由上面可以知道FP16的表示范围在[2^{-24}, 2^15]。一般情况下，网络的梯度值都比较小，如果梯度值小于2^{2^-24}，那么就会发生下溢，数值变为0，导致网络参数无法更新。
其实，我们在训练神经网络的时候，还可以观察到，大部分的梯度值都很小，是无法占满[0, 65504]这个范围的，那么就可以认为的规定小数点左移，使FP16表示的数值范围整体左移即可，例如左移3位，那么表示的数值范围就会整体缩小8倍。具体的做法如下：

• 损失缩放保持小的梯度值
  对较小的梯度进行放缩，使其变大，使原本很小的数值（因为FP16的表示范围导致下溢的）能够表示。具体地，在梯度计算出来之后，将梯度乘上一个常数S，然后进行梯度的反向传播，再进行梯度裁剪或者其他和梯度相关的运算，再将梯度除以常数S，然后进行权重的更新。

  1. Maintain a master copy of weights in FP32
  2. For each iteration:
     a. Make an FP16 copy of the weights
     b. Forward propagation (FP16 weights and activations)
     c. Multiply the resulting loss with the scaling factor_S_
     d. Backward propagation (FP16 weights, activations, and their gradients)
     e. Multiply the weight gradient with 1/_S_
     f. Complete the weight update (including gradient clipping, etc.)

事实上，大多数的能够使用FP32训练的网络，不做任何修改也能够使用FP16进行训练；但是如果网络中涉及一些大的reduce的运算（需要计算mean和variance的Batch-Normalization，SoftMax），这些值的权重需要使用FP32表示。

• 如何选择合适的缩放因子S
  我们也看到按照上述的步骤有一个自由参数S，那么如何选择一个合适的S哪？

  • 选择一个常数
  • 统计梯度绝对值的最大值，该最大值乘以S，不发生FP16的上溢（即不大于65504），那么该S就是合适的

  • 训练中动态调整S的值。
    首先，选择一个相对较大的S开始训练，如果没法发生溢出，经过N次迭代之后，就增大S；如果发生了溢出，就跳过这次权重更新，并且减小S。其实只要不频繁的跳过权重的更新，对于整个网络的训练来说，是机会没有影响的。具体过程可以用如下的语言就行描述：

    1. Maintain a master copy of weights in FP32.
    2. Initialize S to a large value.

    3. For each iteration:
       a. Make an FP16 copy of the weights.
       b. Forward propagation (FP16 weights and activations).
       c. Multiply the resulting loss with the scaling factor S.
       d. Backward propagation (FP16 weights, activations, and their gradients).
       e. If there is an Inf or NaN in weight gradients:

       i. Reduce S.
       ii. Skip the weight update and move to the next iteration.

       f. Multiply the weight gradient with 1/S.
       g. Complete the weight update (including gradient clipping, etc.).
       h. If there hasn’t been an Inf or NaN in the last N iterations, increase S.

---

参考文献

1. http://c.biancheng.net/view/3...
2. Micikevicius, Paulius, et al. "Mixed precision training." arXiv preprint arXiv:1710.03740 (2017).
3. https://mxnet.incubator.apach...
4. https://docs.nvidia.com/deepl...
5. https://mxnet.incubator.apach...