CNN图像分类的小技巧(6): 提升训练效率-混合精度训练

Mixed Precision Training
神经网络模型变得越来越大，所需要的计算和内存资源也越来越多。一个很直觉的想法就是，使用低精度去存储和计算可以有效的降低内存需求和提升运算速度。目前为止大多数模型都是使用单精度(FP32)来存储和计算的，很自然的我们想到是不是可以使用FP16来存储和计算。

1 FP16的好处和问题

FP16相对于FP32来说，只使用16bits来表达浮点数字是FP32的一半，所带来的好处就是：

内存占用更少：在部署的时候模型的大小相比FP32更小，更容易部署到硬件资源有限的嵌入式设备上。在训练的时候模型所占用的内存变小使得可以采用更大的batchsize进行训练。在分布式训练的时候可以减少GPU之间的通信，从而权重可以更快的完成更新继续下一轮的训练。
计算更快：现在不少GPU都会针对FP16的计算进行优化，论文中指出，FP16计算吞吐量可以达到FP32的2到8倍。

但是由于FP16只用了16位比特来存储浮点数，所以FP16能表达的浮点数的范围比FP32要小很多，这会带来两个问题：

数据的溢出：在训练的过程中梯度是个很小的数字而且我们还会乘以一个比较小的学习率，这会导致最终的梯度非常小，当梯度小到比FP16能表达的最小数字还小的时候，FP16只能用0来表示，这样子无法使用梯度去更新权重，模型将无法正常的训练和收敛。
舍入误差：我们举个简单的例子去理解，比如说整数之间的最小间隔是1，当你加一个小于0的数去另外一个整数的时候，这个整数是不会变化的。由于计算机设备的局限性，在计算机的世界里，我们可以理解每个浮点数之间也相隔一个特定的数字，当你加一个数字小于这个特定的数字去另外一个浮点数的时候，那个浮点数就不会更新。所以说当梯度小到比这个特定的数字还小的时候，去更新权重，权重是不会被更新的。

2 混个FP16和FP32一起训练

所以论文提出了一套FP16和FP32混合训练的方法可以解决FP16所带来的问题，从而舍得训练得到的FP16模型有着和FP32模型一样的准确性。这套办法由三个部分组成：

2.1 FP32权重备份

在训练过程中，weights, activations和gradients都用FP16来存储，同时拷贝一份FP32的weights用来更新权重。具体操作是：

首先系统维持一份FP32表达的权重，我们叫master-weights
在需要进行前向和后向训练时，我们拷贝一份FP16格式的master-weights，然后使用它进行前向，后向操作。相比于FP32训练来说，前向后向的运算过程中因为使用FP16使得存储和带宽的需求都减半了。
当需要更新权重时，我们使用后向计算过程中产生的FP16格式的gradient来更新FP32格式的master weights
在接下来的训练中重复步骤2-4.

我们可能会有这样的疑问，在内存中增加一份FP32格式的权重，不是增加了内存的占用嘛？其实在训练过程中，内存主要是被activations占用了，因为大的batchsize和每一层的activations需要保存起来给back-propagation来使用，所以占用了大量的内存，所以但我们使用FP16格式的数据进行前向后向计算时，内存占用最多的地方减半了，所以额外的一份FP32的weights不算什么，总体来说怎个训练过程中的内存占用大概是减半的。

2.2 Loss Scale

这个主要是解决FP16数据的溢出问题，在训练过程中，特别是到了后期和反向传播到很远的时候，梯度的数值很小，很容易超出FP16能表达的最小的值，而变成0. 因此我们可以scale loss，链式反向传播的法则可以确保梯度有着相同的scale，所以在反向传播的过程在我们在做任何跟梯度相关的操作之前我们可以unscale梯度。scale所带来的效果就是将原本不在FP16有效范围内的值平移到FP16的有效范围内。使用scaled梯度进行反向传播，在需要更新权重时将其转化成FP32，然后unsccale，然后进行操作。

2.3 计算精度

神经网络的运算操作基本可以分类三类：vector dot product，reductions，point-wise operations。这三种操作对精度的敏感度都不一样，我们分开介绍：

vector dot product：论文中表明，对于某些模型来说如果想要FP16模型训练维持FP32模型训练一样的精度，需要将FP16向量点积操作过程中的累加(accumulated)操作用FP32来完成。大概的意思是说，用FP16进行乘法和存储，利用FP32来进行加法操作，可以减少加法过程中的舍入误差。
Large reductions：这种运算操作应该使用FP32来进行，同样的使用FP16来进行存储读写，FP32来进行运算操作。这些操作主要发生在batch-normalization层和softmax层，由于这些层的性能瓶颈是带宽，所以使用FP32进行运算并不会拖累训练过程。
Point-wise operation：对于这些操作来说，内存带宽同时限制着它的性能，所以不管使用FP16还是FP32来进行算术运算都不会影响速度。