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

Posted 2023-03-28

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

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

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

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

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

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

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

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