深度学习模型精度与PyTorch模型量化

1. 模型量化是什么

2. Pytorch模型量化

2.1 Tensor的量化

2.2 训练后动态量化Post Training Dynamic Quantization

2.3 训练后静态量化Post Training Static Quantization

2.4 训练时量化Quantization Aware Training

3. 混合精度训练Automatically Mixed Precision

3.1 解决方案

3.2 Pytorch与Tensorflow如何应用混合精度训练

1. 模型量化是什么

PyTorch模型量化官方文档：Quantization — PyTorch 2.0 documentation

当前的深度学习框架大都采用的都是fp32来进行权重参数的存储，比如Python float的类型为双精度浮点数fp64，PyTorch Tensor的默认类型为单精度浮点数fp32。

使用fp32主要存在问题：

模型尺寸大，训练的时候对显卡的显存要求高
模型训练速度慢
模型推理速度慢

PyTorch支持int8量化，使模型大小减少4倍，内存带宽要求减少4倍。与FP32计算相比，硬件对int8计算的支持通常要快2到4倍。量化主要是一种加快推理速度的技术，并且只支持量化运算符的前向传递。简单来说，在深度学习中，量化指的是使用更少的bit来存储原本以浮点数存储的tensor，以及使用更少的bit来完成原本以浮点数完成的计算。

数据类型	取值范围
float16	-65504 ~ 65504
float32	-2^31 ~ 2^31-1
int8	-2^7 ~ 2^7-1 (-128 ~ 127)
uint8	0 ~ 2^8-1 (0~255)

模型量化的好处：

更少的存储开销和带宽需求
更快的计算速度（由于更少的内存访问和更快的int8计算，计算速度快2~4倍）
更低的能耗与占用面积
尚可接受的精度损失（即量化相当于对模型权重引入噪声，所幸CNN本身对噪声不敏感，在模型训练过程中，模拟量化所引入的权重加噪还有利于防止过拟合，在合适的比特数下量化后的模型并不会带来很严重的精度损失）

一个量化后的模型，其部分或者全部的tensor操作会使用int类型来计算，而不是使用量化之前的float类型。当然，量化还需要底层硬件支持，x86 CPU（支持AVX2）、ARM CPU、Google TPU、Nvidia Volta/Turing/Ampere、Qualcomm DSP这些主流硬件都对量化提供了支持。

注：深度学习的模型压缩的主流方法有基于量化的方法、模型剪枝和知识蒸馏，模型量化，这是最广泛使用的模型压缩形式。

2. Pytorch模型量化

PyTorch对量化的支持目前有如下三种方式：

Post Training Dynamic Quantization，模型训练完毕后的动态量化
Post Training Static Quantization，模型训练完毕后的静态量化
QAT（Quantization Aware Training），模型训练中开启量化

2.1 Tensor的量化

>>> x = torch.rand(2,3, dtype=torch.float32) 
>>> x
tensor([[0.6839, 0.4741, 0.7451],
        [0.9301, 0.1742, 0.6835]])

>>> xq = torch.quantize_per_tensor(x, scale = 0.5, zero_point = 8, dtype=torch.quint8)
tensor([[0.5000, 0.5000, 0.5000],
        [1.0000, 0.0000, 0.5000]], size=(2, 3), dtype=torch.quint8,
       quantization_scheme=torch.per_tensor_affine, scale=0.5, zero_point=8)

>>> xq.int_repr()
tensor([[ 9,  9,  9],
        [10,  8,  9]], dtype=torch.uint8)

>>> xdq = xq.dequantize()
>>> xdq
tensor([[0.5000, 0.5000, 0.5000],
        [1.0000, 0.0000, 0.5000]])

quantize_per_tensor函数：使用给定的scale和zp来把一个float tensor转化为quantized tensor
dequantize函数：quantize_per_tensor的反义词，把一个量化tensor转换为float tensor

xdq和x的值已经出现了偏差的事实告诉了我们两个道理：

量化会有精度损失；
选择合适的scale和zp可以有效降低精度损失（例如scale = 0.0036, zero_point = 0）

而在PyTorch中，选择合适的scale和zp的工作就由各种observer来完成。Tensor的量化支持两种模式：per tensor 和 per channel。Per tensor 是说一个tensor里的所有value按照同一种方式去scale和offset； per channel是对于tensor的某一个维度（通常是channel的维度）上的值按照一种方式去scale和offset，也就是一个tensor里有多种不同的scale和offset的方式（组成一个vector），如此以来，在量化的时候相比per tensor的方式会引入更少的错误。PyTorch目前支持conv2d()、conv3d()、linear()的per channel量化。

2.2 训练后动态量化Post Training Dynamic Quantization

这是最简单的量化形式，其中权重被提前量化，而激活在推理过程中被动态量化。这种方法用于模型执行时间由从内存加载权重而不是计算矩阵乘法所支配的情况，对整个模型应用动态量化只需要调用一次torch.quantization.quantize_dynamic()函数即可完成。

torch.quantization.quantize_dynamic(model, qconfig_spec=None, dtype=torch.qint8, mapping=None, inplace=False)

quantize_dynamic这个API把一个float model转换为dynamic quantized model，也就是只有权重被量化的model，dtype参数可以取值 float16 或者 qint8。当对整个模型进行转换时，默认只对以下的op进行转换：

Linear
LSTM
LSTMCell
RNNCell
GRUCell

为啥呢？因为dynamic quantization只是把权重参数进行量化，而这些layer一般参数数量很大，在整个模型中参数量占比极高，因此边际效益高。对其它layer进行dynamic quantization几乎没有实际的意义。

对于一个默认行为下的quantize_dynamic调用，举例如下：

#原始网络
Net(
  (conv): Conv2d(1, 1, kernel_size=(1, 1), stride=(1, 1), bias=False)
  (fc): Linear(in_features=3, out_features=2, bias=False)
  (relu): ReLU()
)

#quantize_dynamic 后
Net(
  (conv): Conv2d(1, 1, kernel_size=(1, 1), stride=(1, 1), bias=False)
  (fc): DynamicQuantizedLinear(in_features=3, out_features=2, dtype=torch.qint8, qscheme=torch.per_tensor_affine)
  (relu): ReLU()
)

可以看到，除了Linear，其它op都没有变动。而Linear被转换成了DynamicQuantizedLinear，DynamicQuantizedLinear就是torch.nn.quantized.dynamic.modules.linear.Linear类。没错，quantize_dynamic API的本质就是检索模型中op的type，如果某个op的type属于字典DEFAULT_DYNAMIC_QUANT_MODULE_MAPPINGS的key，那么，这个op将被替换为key对应的value：

# Default map for swapping dynamic modules
DEFAULT_DYNAMIC_QUANT_MODULE_MAPPINGS = {
    nn.GRUCell: nnqd.GRUCell,
    nn.Linear: nnqd.Linear,
    nn.LSTM: nnqd.LSTM,
    nn.LSTMCell: nnqd.LSTMCell,
    nn.RNNCell: nnqd.RNNCell,
}

总结：Post Training Dynamic Quantization，简称为Dynamic Quantization动态量化，或者叫作Weight-only的量化，是提前把模型中某些op的参数量化为INT8，然后在运行的时候动态的把输入量化为INT8，然后在当前op输出的时候再把结果requantization回到float32类型。动态量化默认只适用于Linear以及RNN的变种。

2.3 训练后静态量化Post Training Static Quantization

这是最常用的量化形式，其中权重是提前量化的，并且基于在校准过程中观察模型的行为来预先计算激活张量的比例因子和偏差。CNN是一个典型的用例，训练后量化通常是在内存带宽和计算节省都很重要的情况下进行的。

与dynamic quantization的相同点和区别：

相同点就是，都是把网络的权重参数转从float32转换为int8
不同点是，需要把训练集或者和训练集分布类似的数据喂给模型（注意没有反向传播），然后通过每个op输入的分布特点来计算activation的量化参数（scale和zp）——称之为Calibrate（定标）。

静态量化包含有activation了，即post process，也就是op forward之后的后处理。为什么静态量化需要activation呢？因为静态量化的前向推理过程自(始+1)至(终-1)都是INT计算，activation需要确保一个op的输入符合下一个op的输入。

进行训练后量化的一般过程如下所示：

步骤1-准备模型：通过添加QuantStub和DeQuantStub模块，指定在何处显式量化和反量化激活值；确保不重复使用模块；将需要重新量化的任何操作转换为模块的模式；
步骤2-将诸如conv + relu或conv + batchnorm + relu之类的组合操作融合在一起，以提高模型的准确性和性能；
步骤3-指定量化方法的配置，例如选择对称或非对称量化以及MinMax或L2Norm校准技术；
步骤4- 插入torch.quantization.prepare()模块来在校准期间观察激活张量；
步骤5-使用校准数据集对模型执行校准操作；
步骤6-使用torch.quantization.convert() 模块来转化模型，具体包括计算并存储每个激活张量要使用的比例和偏差值，并替换关键算子的量化实现等。

具体参考：PyTorch的量化 - 知乎

动态量化和静态量化的最大区别：

静态量化的float输入必经QuantStub变为int，此后到输出之前都是int
动态量化的float输入是经动态计算的scale和zp量化为int，op输出时转换回float

2.4 训练时量化Quantization Aware Training

在极少数情况下，训练后量化不能提供足够的准确性，可以插入torch.quantization.FakeQuantize()模块执行训练时量化。计算将在FP32中进行，但将值取整并四舍五入以模拟INT8的量化效果。具体的量化步骤如下所示：

步骤1-准备模型：通过添加QuantStub和DeQuantStub模块，指定在何处显式量化和反量化激活值；确保不重复使用模块；将需要重新量化的任何操作转换为模块的模式；
步骤2-将诸如conv + relu或conv + batchnorm + relu之类的组合操作融合在一起，以提高模型的准确性和性能；
步骤3-指定伪量化方法的配置，例如选择对称或非对称量化以及MinMax或L2Norm校准技术；
步骤4-插入torch.quantization.prepare_qat() 模块，该模块用来在训练过程中的模拟量化；
步骤5-训练或者微调模型；
步骤6-使用torch.quantization.convert() 模块来转化模型，具体包括计算并存储每个激活张量要使用的比例和偏差值，并替换关键算子的量化实现等。

具体参考：PyTorch的量化 - 知乎

QAT总结：

# 原始的模型，所有的tensor和计算都是浮点
previous_layer_fp32 -- linear_fp32 -- activation_fp32 -- next_layer_fp32
                      /
    linear_weight_fp32

# 训练过程中，fake_quants发挥作用
previous_layer_fp32 -- fq -- linear_fp32 -- activation_fp32 -- fq -- next_layer_fp32
                           /
   linear_weight_fp32 -- fq

# 量化后的模型进行推理，权重和输入都是int8
previous_layer_int8 -- linear_with_activation_int8 -- next_layer_int8
                     /
   linear_weight_int8

3. 混合精度训练Automatically Mixed Precision

使用低精度计算对模型进行优化，推理过程中，模型优化目前比较成熟的方案就是fp16量化和int8量化；训练方面的方案是混合精度训练，它的基本思想很简单:，精度减半(fp32→ fp16) ，训练时间减半。与单精度浮点数float32（32bit，4个字节）相比，半精度浮点数float16仅有16bit，2个字节组成。

在训练过程中将 FP32 替代为 FP16，有两个好处：

减少显存占用：FP16 的显存占用只有 FP32 的一半，可以用更大的 batch size
加速训练：使用 FP16，模型的训练速度几乎可以提升 1 倍

训练过程中，直接使用半精度进行计算会导致的两个问题：

舍入误差(Rounding Error)：对足够小的浮点数执行的任何操作都会将该值四舍五入到零，在反向传播中很多甚至大多数梯度更新值都非常小，在反向传播中舍入误差累积可以把这些数字变成0或者nan，这会导致不准确的梯度更新，影响网络的收敛
溢出错误(Grad Overflow / Underflow)：精度下降(小数点后16相比较小数点后8位要精确的多)会导致得到的值大于或者小于fp16的有效动态范围，也就是上溢出或者下溢出

混合精度 (Automatically Mixed Precision, AMP) 仅仅只需几行代码，就能让显存占用减半，训练速度加倍。 AMP 技术是由百度和 NIVDIA 团队在 2017 年提出的 (Mixed Precision Training)，该成果发表在 ICLR 上。PyTorch 1.6之前，大家都是用 NVIDIA 的 apex 库来实现 AMP 训练。1.6 版本之后，PyTorch 出厂自带 AMP。

混合精度训练是在尽可能减少精度损失的情况下利用半精度浮点数加速训练，它使用FP16即半精度浮点数存储权重和梯度，在减少占用内存的同时起到了加速训练的效果。用FP16代替原FP32神经网络计算的最大问题就是精度损失。

3.1 解决方案

1）为每个权重保留一份FP32的副本

为了实现 FP16 的训练，我们需要把模型权重和输入数据都转成 FP16，反向传播的时候就会得到 FP16 的梯度。如果此时直接进行更新，因为梯度 * 学习率的值往往较小，和模型权重的差距会很大，可能会出现舍入误差的问题。

所以解决思路是：将模型权重、激活值、梯度等数据用 FP16 来存储，同时维护一份 FP32 的模型权重副本用于更新。在反向传播得到 FP16 的梯度以后，将其转化成 FP32 并 unscale，最后更新 FP32 的模型权重。因为整个更新过程是在 FP32 的环境中进行的，所以不会出现舍入误差。

FP32 权重备份解决了反向传播的舍入误差问题。

2）Loss-scaling

为了解决下溢出的问题对计算出来的 loss 值进行缩放 (scale)，由于链式法则的存在，对 loss 的缩放会作用在每个梯度上。缩放后的梯度，就会平移到 FP16 的有效范围内。这样就可以用 FP16 存储梯度而又不会溢出了。此外，在进行更新之前，需要先将缩放后的梯度转化为 FP32，再将梯度反缩放 (unscale) 回去。

缩放因子 (loss_scale) 一般都是框架自动确定的，只要没有发生 inf 或者 nan，loss_scale 越大越好。因为随着训练的进行，网络的梯度会越来越小，更大的 loss_scale 可以更加充分地利用 FP16 的表示范围。

3）改进算数方法：FP16 * FP16 + FP32

对于那些在 FP16 环境中运行不稳定的模块，强制它在 FP32 的精度下运行。比如需要计算 batch 均值的 BN 层就应该在 FP32 下运行，否则会发生舍入误差。以 BN 层为例，将其权重转为 FP32，并且将输入从 FP16 转成 FP32，这样就可以保证整个模块是在 FP32 下运行的。

3.2 Pytorch与Tensorflow如何应用混合精度训练

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'

Reference：