节省显存新思路，在 PyTorch 里使用 2 bit 激活压缩训练神经网络( 二 )

把fp32浮点数量化为2-bit整数是一个有损压缩，会引入一些误差。论文从理论上分析了量化引入的误差是如何影响训练的收敛性的。
第一，存在一个随机化的量化策略，使得使用有损量化压缩后，估计出的有损梯度是原梯度的一个无偏估计。

文章图片
在这一条件下，我们套用已有的随机梯度下降收敛性定理，得出最后收敛时的误差会被梯度的方差所限制。
第二，我们推导出了使用量化压缩之后，随机梯度下降计算出的梯度的方差。

文章图片
等号右边的第一项是随机梯度下降在minibatch采样时产生的方差，等号右边的第二项是有损压缩额外引入的方差。这条公式显示地刻画了有损压缩带来的影响。注意到，当有损量化压缩带来的方差远小于原来随机梯度下降自带的方差时， ActNN引入的有损压缩就不会影响训练的收敛性。更多关于公式的推导和可视化参见文末的论文链接。论文对不同的算子（conv2d ， batchnorm ， linear等）都提供了详细的分析。
由上述公式启发，我们提出了一些新的量化技巧用于降低有损压缩引入的额外方差。我们引入了新的量化技巧（Per-groupQuantization ， Fine-GrainedMixed-Precision ， RuntimeAdaptation）来利用梯度在不同样本，不同纬度，不同层之间的异构特性。最后的压缩算法会分配更多的bit给更重要的激活值。平均每个浮点数分配到2bit 。
在具体实现压缩算法时，还有很多可以调节的参数。这里产生了一个内存节省和训练速度的取舍。一般来说，使用更复杂的压缩算法可以节省更多的内存，但是也会引入更多额外的开销，使训练速度变慢。为了给用户较大的灵活性， ActNN提供了5个优化等级L1-L5供用户选择。低的优化等级节省的内存较少，但是运行速度快。高的优化等级节省的内存多，但是运行也更慢。在最高优化等级L5下， ActNN会结合一个简单的内存交换策略，将压缩后的激活值移到CPU内存上，进一步节省内存。
实现
要在PyTorch实现ActNN算法非常简单。对于一个PyTorchnnModule ，我们只需要在其forward函数里加入量化压缩，在其backward函数里加入解压缩操作。所有的计算还是在fp32下进行，与原来一样，伪代码如下图所示。
ActNN为大部分常用的PyTorchnn.Module实现了使用量化压缩的版本。用户只需将模型里的所有PyTorchnn.Module替换成ActNN对应的Module（如把nn.Conv2d替换成actnn.Conv2d) ，即可节省内存，不需要更改其他代码。 ActNN同时也提供了一个wrapper实现一行代码自动替换。

文章图片
实验结果
因为ActNN进行的是有损压缩，所以最重要的一点是先验证ActNN是否会影响模型的精度。下图是使用ActNN在ImageNet上训练ResNet-50的结果。 FP代表普通的fp32训练， BLPA是来自NeurIPS2019的一个相关工作。可以看到，在ActNN的2-bit压缩模式下，模型几乎没有损失精度。在更极限的1.25bit的情况下， ActNN也能收敛，只不过会损失一些精度。而之前的工作BLPA在小于4bit的情况就下无法收敛。

文章图片
我们还在图像分割，物体检测，以及自监督学习等多个任务上进行了实验。 ActNN都能在2-bit压缩模式下达到和普通fp32几乎一样的结果。在部分任务上，因为ActNN可以使用更大的batchsize ，甚至可以取得更好的测试结果。详细的实验结果和训练记录参见文末的论文与github链接。