【DL】神经网络混合精度训练

以前介绍过了神经网络分布式训练，没怎么研究混合精度，觉得就是都转成FP16就行了，最近才发现仍是有些东西的，因此看了下百度和英伟达合做的MIXED PRECISION TRAINING，把细节记下来。

1. 混合精度训练

混合精度训练是在尽量减小精度损失的状况下利用半精度浮点数加速训练。它使用FP16即半精度浮点数存储权重和梯度。在减小占用内存的同时起到了加速训练的效果。

IEEE标准中的FP16格式以下：

取值范围是5.96× 10−8 ~ 65504，而FP32则是1.4×10-45 ~ 3.4×1038。

从FP16的范围能够看出，用FP16代替原FP32神经网络计算的最大问题就是精度损失。

文中提出了三种避免损失的方法：

1.1. 为每一个权重保留一份FP32的副本

在前向和反向时使用FP16，整个过程变成：权重从FP32转成FP16进行前向计算，获得loss以后，用FP16计算梯度，再转成FP32更新到FP32的权重上。这里注意获得的loss也是FP32，由于涉及到累加计算（参见下文）。

用FP32保存权重主要是为了不溢出，FP16没法表示2e-24如下的值，一种是梯度的更新值过小，FP16直接变为了0；二是FP16表示权重的话，和梯度的计算结果也有可能变成0。实验代表，用FP16保存权重会形成80%的精度损失。

1.2. Loss-scaling

放大FP16的梯度，变成FP32的梯度，更新时缩放回来。下图能够看到，不少激活值比较小，没法用FP16表示。所以在前向传播后对loss进行扩大（固定值或动态值），这样在反响传播时全部的值也都扩大了相同的倍数。在更新FP32的权重以前unscale回去。

在训练期间，一些权重梯度指数很小，在FP16格式下会变成0。为了解决这个问题，在反向传播开始时引入一个换算系数，来缩放损失，梯度也经过链式法则同时逐渐扩大，并能在FP16中表示出来。

1.3. 改进算数方法：FP16 * FP16 + FP32。

通过实验，做者发现将FP16的矩阵相乘后和FP32的矩阵进行加法运算，写入内存时再转回FP16能够得到较好的精度。英伟达V系列GPU卡中的Tensor Core（上图）也很支持这种操做。所以，在进行大型累加时（batch-norm、softmax），为防止溢出都须要用FP32进行计算，且加法主要被内存带宽限制，对运算速度不敏感，所以不会下降训练速度。另外，在进行Point-wise乘法时，用FP16或者FP32均可以，引用原文感觉如下：

Point-wise operations, such as non-linearities and element-wise matrix products, are memory- bandwidth limited. Since arithmetic precision does not impact the speed of these operations, either FP16 or FP32 math can be used.

2. 实验结果

从下图的Accuracy结果能够看到，混合精度基本没有精度损失：

Loss scale的效果：

3. 如何应用MP

Pytorch可使用英伟达的开源框架APEX，支持混合进度和分布式训练：

model, optimizer = amp.initialize(model, optimizer, opt_level="O1")
with amp.scale_loss(loss, optimizer) as scaled_loss:
    scaled_loss.backward()复制代码

Tensorflow就更简单了，已经有官方支持，只须要训练前加一句：

export TF_ENABLE_AUTO_MIXED_PRECISION=1
# 或者
import os
os.environ['TF_ENABLE_AUTO_MIXED_PRECISION'] = '1'复制代码

参考：Automatic Mixed Precision for Deep Learning