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

背景

随着深度神经网络的快速发展，网络结构和数据集呈现复杂化，巨大化；如何快速的训练一个深度神经网络面临着严峻的挑战。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内部原理

未完待续

---

参考文献

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...