深度学习中模型量化-入门篇

一、模型量化

1.1 什么是量化？

模型量化 | 455のblog (jilei-hou.github.io)

简单来说，模型量化（Model Quantization）就是通过某种方法将浮点模型转为定点模型。比如说原来的模型里面的权重（weight）都是float32，通过模型量化，将模型变成权重（weight）都是int8的定点模型。

随着深度学习（Deep Learning）的发现，其在计算机视觉（Computer Vision， CV）和自然语言处理（Natural Language Processing，NLP）等领域都取得了巨大的成功。通过深度学习，我们可以得到用于处理各种任务的高性能模型，这些模型大多都很复杂、一般只适合在GPU上进行推理，并不适合在板端进行推理，然而在实际应用时，很多场景都需要将模型部署到板端。

为了解决模型难以部署到板端的问题，我们就需要通过模型量化来降低模型的复杂性，这个过程不可避免的会发生精度损失。

量化前的浮点模型和量化后的定点模型的特点：

量化前的浮点模型	量化后的定点模型
参数量大（float32）	压缩参数（int8）
计算量大	提升速度
内存占用多	内存占用少
精度高	精度损失

模型量化就是建立一种浮点数据和定点数据间的映射关系，使得以较小的精度损失代价获得了较大的收益，要弄懂模型量化的原理就是要弄懂这种数据映射关系。

1.1.1 pytorch支持的量化模式

1.1.2 PTQ 与 QAT

• 后训练量化（Post-Training Quantization, PTQ）

后训练量化(PTQ)可以在没有原始的训练过程的情况下，就能将预训练的FP32模型直接转换为定点模型网络。PTQ最大的特点就是不需要数据或者只需要很少的校准数据集。且PTQ几乎不需要调整超参数，使得我们可以很方便的进行模型量化

后训练量化是比较常用且方便快捷的一种方式，需要我们事先将模型训练好，然后加载权重进行量化。

• 量化感知训练（Quantization Aware Training, QAT）

模型量化过程其实就是在做一件事，就是找阈值或者scale。

在PTQ中，阈值或者scale一般是通过统计的方法，然后人工通过一些分布相似性得到的，然而，这肯定是有误差的。而且，由于量化是每层独立进行的，所以每层的量化是不依赖于前一层量化的结果的，这就导致了在实际的inference过程中会出现误差累积的情况，进一步影响量化后的性能。所以，我们需要一种可学习的scale。QAT就是在做这样一件事情。

简单概括就是，我们在网络训练过程去模拟量化，我们通过设定一个可学习的scale，这个scale一般可以与权重或者激活值相绑定，然后我们利用一个量化过程 q = round(r/s)*127，将需要量化的值量化到0-127之间，再接着一个反量化过程q * s，就实现了一个误差的传递，接着我们利用反量化后的结果继续前传，最后得到loss，我们求量化后权重的梯度，并用它来更新量化前的权重，使得这种误差被网络抹平，让网络越来越像量化后的权重靠近，最后我们得到了量化后的缩放因子s。而这一系列操作都可以写成网络中的一个op，实现网络的正常训练。

1.2 线性量化与非线性量化

1.2.1 线性量化

采用相同的量化间隔对输入作量化

根据Z (zero_point) 是否为0，线性量化可以分为对称量化和非对称量化

1.2.2 非对称量化

浮点0对应的值不是定点0

浮点和定点之间的映射公式：

Q = clamp(Round(R/S + Z)) = Qmax, R∈float且R>Tmax

Q = clamp(Round(R/S + Z)) = Round(R/S + Z), R∈float且Tmin<R<Tmax

Q = clamp(Round(R/S + Z)) = Qmin, R∈float且R<Tmin

R = (Q - Z) * S

其中，Q表示量化后的定点数，R表示量化前的浮点数，Z就是zero_point，即浮点数映射到定点之后，浮点0所对应的定点值。S就是scale，即缩放尺度。Round()函数就是四舍五入。clamp()函数的作用是把一个值限制在一个上限和下限之间。Tmax表示浮点数的最大阈值，Tmin表示浮点数的最小阈值。Qmax表示定数的最大值，Qmin表示定点数的最小值。

通过换算可以得到阈值和线性映射参数 S 和 Z 的数学关系，在确定了阈值后，也就确定了线性映射的参数。

S = (Tmax - Tmin) / (Qmax - Qmin)

Z = Qmax - Tmax/S

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

从上述的映射关系中，如果知道了阈值，那么其对应的线性映射参数也就知道了，整个量化过程也就明确了。

那么该如何确定阈值呢？

一般来说，对于权重的量化，由于权重的数据分布是静态的，一般直接找出 MIN 和 MAX 线性映射即可；而对于推理激活值来说，其数据分布是动态的，为了得到激活值的数据分布，往往需要一个所谓校准集的东西来进行抽样分布，有了抽样分布后再通过一些量化算法进行量化阈值的选取（饱和量化）。

1.2.3 量化（一种特殊的非对称量化）

浮点0对应的值就是定点0。对称量化对于正负数不均匀分布的情况不够友好，比如如果浮点数全部是正数，量化后的数据范围是[0, 127], [-128, 0]的范围就浪费了，减弱了int8数据的表示范围

1.2.4 非线性量化

对输入进行量化时，大的输入采用大的量化间隔，小的输入采用小的量化间隔。

1.3 量化举例

举例：

模型训练后权重或激活值往往在一个有限的范围内分布，如权重值范围为[-2.0, 6.0]，即Tmax = 6.0，Tmin = -2.0（非饱和量化）。然后我们用int8进行模型量化，则定点量化值范围为[-128, 127]，即Qmax = 127，Qmin = -127，那么S和Z的求值过程如下：

S = 6.0 - (-2.0) / (127 - (-128)) = 8.0 / 255 ≈ 0.03137255

Z = 127 - 6.0 / 0.03137255 ≈ 127 - 191.25 ≈ -64.25 ≈ -64

可以得到如下对应关系：

浮点数	定点数
6.0	-128
0	-64
-2.0	127

得到量化参数S和Z后，我们就可以求任意一个浮点数对应的定点数，比如说有一个权重等于0.28，即R=0.28

Q = 0.28 / 0.03137255 + (-64) ≈ -55

反量化：

根据 *R = (Q-Z)S = （（-55）- （-64） ）*0.0313 = 0.28

二、 量化训练的流程

1. 设置config

   qconfig 需要设置到模型或者模型的子module上，是QConfig的一个实例

   Qconfig封装了两个observer，一个是激活函数，一个是权重的

   • 如果要部署在x86服务器上

   backend = "fbgemm" #量化操作符在x86机器上通过FBGEMM后端运行 部署在ARM上时，可以使用qnnpack后端
   torch.backends.quantized.engine = backend
   yolox_model.qconfig = torch.quantization.get_default_qconfig(backend)  # 不同平台不同配置
   torch.quantization.prepare(yolox_model, inplace=True) #准备量化
   

   • 如果要部署在ARM（进阶精简指令集机器）上

   backend = "qnnpack"
   

   observer根据上述的四元组（TMIN,TMAX,QMIN,QMAX），分别表示输入的权重（浮点数）数据的最大值最小值分布，和量化后的（定点数）的最大值和最小值分布。

2. fuse_model

   合并可以合并的一些层

   def fuse_model(self):
   	for m in self.modules():
   		if isinstance(m, BaseConv):
   			torch.quantization.fuse_modules(m, ['conv', 'bn', ], inplace=True)
   

   注意：

   在深度学习模型量化过程中，一些层是可以合并的。通常情况下，一个层由若干个神经元组成，它们组成一个大的节点或模块，同时使用softmax或多层激活函数（如1x1、3x3 等）对输入进行激活。以下是一些可以合并的层：

   1. 卷积层和全连接层可以合并成卷积-池化层，能够加速神经网络的训练，但容易过拟合。
   2. 一些线性回归层（如普通最小二乘、Ridge、Leaky ReLU 等）和softmax层可以合并成一个线性回归层，这样可以减少网络结构的复杂度，提高训练效果。
   3. 在一些任务中，可以将一些线性层（如ReLU、tanh 等）和sigmoid 激活函数（如sigmoid 或sigmoid-tanh 函数）合并成一个线性层，这样可以减少网络结构的复杂度，提高训练效果。
   4. 在一些深度残差网络（Deep Residual Learning）中，可以将多个卷积层和上采样层（upsampling layer）合并成一个大的卷积层和上采样层，这样可以加速网络的训练和提高网络的性能。

   总之，在深度学习模型量化过程中，可以根据任务需求和网络结构来选择合适的层进行合并。
   

3. prepare

    torch.quantization.prepare(yolox_model, inplace=True) #准备量化
   

   用来给每个子module插入observer，收集和定标数据，通过观察数据得到四元组中的值，然后在第五步中可以计算得到Scale和zero_point这两个参数的值

4. feed data

   没有反向传播，获取数据分布的特点。

   对于fbgemm

   • 激活函数

     使用Histogram Observer（直方图观察器）是一种用于估计权重分布的观察器类型。它在量化过程中收集权重值的直方图统计信息，并基于这些统计信息计算量化参数。由输入数据+权重值根据L2Norm确定

   • 权重

     针对权重的观察器，默认情况下，常见的做法是使用Per-channel Observer（逐通道观察器）。

     Per-channel Observer意味着对于具有多个通道的权重张量，观察器会针对每个通道分别估计观察统计信息。这样做的目的是为了更好地适应权重在不同通道上的分布特征，提高量化的精度和表现。

     通过逐通道观察器，可以计算每个通道的最小值和最大值，并在量化过程中使用这些通道-wise 的统计信息来确定相应通道的量化参数。这样可以更准确地量化权重，并提高模型在量化后的准确性。

5. convert model

   第4步完成后，各个权重的四元组、各个激活函数的最大值最小值也确定了，可以使用convertAPI

   torch.quantization.convert(yolox_model, inplace=True)
   

   convert的过程中会调用from_float()api，这个api会使用前面的四元组信息计算出weight和activation的Scale和zero_point，用于量化。

   import torch
   import numpy as np
   from PIL import Image
   
   import cv2, copy, os
   from nets.yolo_quant import YoloBody
   import torch.nn.quantized.modules.activation
   from utils.utils import cvtColor, get_classes, preprocess_input, resize_image
   import torch.nn.quantized
   # import tvm.relay.frontend.qnn_torch
   
   def feed_data(epoch, img_num):
       '''
          epoch:要喂数据的轮数
          img_num:每轮要喂的图片的数量
       '''
       yolox_model = YoloBody(20, 's').eval() # 是YoloBody类的实例，调用eval()方法，将模型设置为评估模式
       weight_path = 'logs/ep100-loss4.387-val_loss4.383.pth'  # 更换为训练出来的浮点权重
       weight = torch.load(weight_path, map_location='cpu')
       yolox_model.load_state_dict(weight)
       yolox_model.eval().fuse_model() # 实例调用fuse_model()方法，将卷积层和BN层融合，加快模型推理速度
       # 设置量化参数
       backend = "fbgemm" #量化操作符在x86机器上通过FBGEMM后端运行 部署在ARM上时，可以使用qnnpack后端
       torch.backends.quantized.engine = backend
   
       yolox_model.qconfig = torch.quantization.get_default_qconfig(backend)  # 不同平台不同配置
       torch.quantization.prepare(yolox_model, inplace=True) #准备量化
   
       image_path = 'VOCdevkit/VOC2007/JPEGImages/'  # 图片路径更换为数据集路径能够读取其中的图片进行下面的前向预测过程
       img_list = os.listdir(image_path)
       for i in range(0, epoch):
           for iter, img_name in enumerate(img_list):
               yolox_model.train()
               image = Image.open(image_path + img_name)
               crop_image = image.convert('L')
               # 对每张图像，首先将其打开为Image对象，然后进行灰度转换和尺寸调整操作
               crop_image = crop_image.resize((640, 640), Image.BICUBIC)
               photo = preprocess_input(np.array(crop_image, dtype=np.float32))
               image_data = np.expand_dims(photo, axis=0)
               image_data = np.expand_dims(image_data, axis=0)
   
               with torch.no_grad():
                   images = torch.from_numpy(np.asarray(image_data))
                   # print(type(images))
                   outputs = yolox_model(images)
               print('{}--{} finish'.format(iter + 1, img_name))
               if (iter + 1) == img_num:
                   break
   
       torch.quantization.convert(yolox_model, inplace=True)
       torch.jit.save(torch.jit.script(yolox_model.eval()),
                      "quant_pth_test/yolox_quantization_weight_epoch{}.pth".format(epoch + 1))
        
       # torch.save(yolox_model.state_dict(), "quant_pth_68/yolox_para_10.pth".format(epoch + 1))
   
       print("=========================quantized========================")
   
       for k, v in yolox_model.state_dict().items():
       # 打印量化结果。遍历模型的状态字典，保存权重相关的量化参数（如缩放因子、零点等）和权重的浮点表示形式
           if 'weight' in k:   #权重相关的量化参数
               np.save('./para_test/' + k + '.scale', v.q_per_channel_scales()) # 保存缩放因子，不同通道的缩放因子不同
               np.save('./para_test/' + k + '.zero_point', v.q_per_channel_zero_points()) # 保存零点，不同通道的权重在量化过程中的零偏移量不同
               np.save('./para_test/' + k + '.int', v.int_repr()) # 保存权重的整数表示形式，将浮点权重乘以缩放因子，然后加上零点，再四舍五入取整
               np.save('./para_test/' + k, v.dequantize().numpy()) # 保存权重的浮点表示形式，还原量化后的权重
           elif 'bias' in k:  # 偏置相关的量化参数
               np.save('./para_test/' + k, v.detach().numpy())
           elif 'zero_point' in k: # 零点相关的量化参数
               np.save('./para_test/' + k, v.detach().numpy())
           elif 'scale' in k: # 缩放因子相关的量化参数
               np.save('./para_test/' + k, v.detach().numpy())
   
   feed_data(1, 2)
   

2.1 简单模型实战

下面我们定义一些简单的层次来实现量化，主要是融合一些特定层

import torch
import torch.nn as nn
import torch.quantization
from torch.quantization import QuantStub, DeQuantStub, fuse_modules
class ConvNet(nn.Module):
    def __init__(self):
        super(ConvNet, self).__init__()
        self.conv1 = nn.Conv2d(3, 16, kernel_size=3, stride=1, padding=1)
        self.bn1 = nn.BatchNorm2d(16)
        self.relu1 = nn.ReLU(inplace=True)
        self.conv2 = nn.Conv2d(16, 32, kernel_size=3, stride=1, padding=1)
        self.bn2 = nn.BatchNorm2d(32)
        self.relu2 = nn.ReLU(inplace=True)
        self.fc = nn.Linear(32 * 8 * 8 * 16, 10)
        self.quant = QuantStub()
        self.dequant = DeQuantStub()

    def forward(self, x):
        print(x.dtype)
        x = self.quant(x) # 量化
        # 查看类型
        print(x.dtype) 
        x = self.relu1(self.bn1(self.conv1(x)))
        x = self.relu2(self.bn2(self.conv2(x)))
        x = x.reshape(x.size(0), -1)
        x = self.fc(x)
        return self.dequant(x)  # 反量化

    # 聚合
    def fuse_model(self):
        conv_name_dict = {}
        bn_name_dict = {}
        for idx, tu in enumerate(self.named_modules()):
            module_name, module = tu
            if isinstance(module, nn.Conv2d):
                conv_name_dict[idx] = module_name
            elif isinstance(module, nn.BatchNorm2d):
                bn_name_dict[idx] = module_name

        for idx, bn_name in bn_name_dict.items():
            assert (idx - 1) in conv_name_dict.keys()
            conv_name = conv_name_dict[idx - 1]
            fuse_modules(self, [conv_name, bn_name], inplace=True)

# 创建一个实例
model = ConvNet()
print(model)


## 量化过程

input_tensor = torch.randn(1, 3, 32, 32)

# 进入测试模式
model.eval().fuse_model()  # 调用fuse_model()函数，将BN层和卷积层融合

# 设置量化参数
backend = 'fbgemm'  # 量化后端
torch.backends.quantized.engine = backend
model.qconfig = torch.quantization.get_default_qconfig(backend)
torch.quantization.prepare(model, inplace=True)  # 准备量化
# convert
torch.quantization.convert(model, inplace=True)
# 前向预测
output = model(input_tensor)

print(model)

量化前的网络结构如下：

量化过程中我们可以看到数据类型发生了变化

当我们融合conv2d和bn两层之后

可以看到模型融合成功，原始网络中的conv会被替换为新的合并后的module,bn会被替换成Identity()，可以理解为占位符。实际上，上述结构中，还可以进一步融合relu，我们点开fuse_modules包，可以看到pytorch可以量化融合的层：

当然我们也可以参考官方的手册说明：Pytorch Mobile Performance Recipes — PyTorch Tutorials 2.0.1+cu117 documentation

参考资料

1. 一文搞懂模型量化 - 知乎 (zhihu.com)
2. Pytorch Mobile Performance Recipes — PyTorch Tutorials 2.0.1+cu117 documentation