Pytorch量化之Post Train Static Quantization(训练后静态量化)

使用Pytorch训练出的模型权重为fp32，部署时，为了加快速度，一般会将模型量化至int8。与fp32相比，int8模型的大小为原来的1/4, 速度为2~4倍。
Pytorch支持三种量化方式：

• 动态量化（Dynamic Quantization）: 只量化权重，激活在推理过程中进行量化
• 静态量化（Static Quantization）: 量化权重和激活
• 量化感知训练（Quantization Aware Training，QAT）: 插入量化算子后进行训练，主要在静态量化精度不满足需求时进行。
  大多数情况下，我们只需要进行静态量化，少数情况下在量化感知训练不满足时使用QAT进行微调。所以本篇只重点讲静态量化，并且理论部分先略过（后面再专门总结），只关注实操。
  注：下面的代码是在pytorch1.10下，后面Pytorch对量化的接口有调整
  官方文档：Quantization — PyTorch 1.10 documentation

动态模式（Eager Mode）与静态模式（fx graph）

Pytorch支持用2种方式量化，一种是动态图模式，也是我们日常使用Pytorch训练所使用的方式，使用这种方式量化需要自己手动修改网络结构，在支持量化的算子前、后插入量化节点，优点是方便调试。静态模式则是由pytorch自动在计算图中插入量化节点，不需要手动修改网络。
网络上大部分的教程都是基于静态模式，这种方式比较大的问题就是需要手动修改网络结构，官方教程里的网络是属于demo型, 其中的QuantStub和DeQuantStub就分别是量化和反量化的节点：

# define a floating point model where some layers could be statically quantized
class M(torch.nn.Module):
    def __init__(self):
        super(M, self).__init__()
        # QuantStub converts tensors from floating point to quantized
        self.quant = torch.quantization.QuantStub()
        self.conv = torch.nn.Conv2d(1, 1, 1)
        self.relu = torch.nn.ReLU()
        # DeQuantStub converts tensors from quantized to floating point
        self.dequant = torch.quantization.DeQuantStub()

    def forward(self, x):
        # manually specify where tensors will be converted from floating
        # point to quantized in the quantized model
        x = self.quant(x)
        x = self.conv(x)
        x = self.relu(x)
        # manually specify where tensors will be converted from quantized
        # to floating point in the quantized model
        x = self.dequant(x)
        return x

Pytorch对于很多网络层是不支持量化的（比如很常用的Prelu），如果我们用这种方式，我们就必须在这些不支持的层前面插入DeQuantStub，然后在支持的层前面插入QuantStub。笔者体验下来，体验很差，个人觉得不太实用，会破坏原来的网络结构。
而静态图模式，我们只需要调用Pytorch提供的接口将原模型转换一下即可，不需要修改原来的网络结构文件，个人认为实用性更强。

静态模式量化

1. 载入fp32模型，并转成fx graph

其中量化参数有‘fbgemm’和‘qnnpack’两种，前者在x86运行，后者在arm运行。

model_fp32 = torch.load(xxx)
model_fp32_quantize = copy.deepcopy(model_fp32)
qconfig_dict = {"": torch.quantization.get_default_qconfig('fbgemm')}
model_fp32_quantize.eval()
# prepare

model_prepared = quantize_fx.prepare_fx(model_fp32_quantize, qconfig_dict)
model_prepared.eval()

2.读取量化数据，标定（Calibration）量化参数

标定的过程就是使用模型推理量化图片，然后统计权重和激活分布，从而得到量化参数。量化图片一般来源于训练集（几百张左右，根据测试情况调整）。量化图片可以通过Pytorch的Dataloader读取，也可以直接自行实现读图片然后送入网络。

### 使用dataloader读取
for i, (data, label) in enumerate(train_loader):
    data = data.to(torch.device("cpu:0"))
    outputs = model_prepared(data)
    print("calibrating {}".format(i))
    if i > 1000:
        break

3. 转换为量化模型并保存

quantized_model = quantize_fx.convert_fx(model_prepared)
torch.jit.save(torch.jit.script(quantized_model), "quantized_model.pt")

速度测试

量化后的模型使用方法与fp32模型一样：

import torch
import cv2
import numpy as np
torch.set_num_threads(1)

fused_model = torch.jit.load("jit_model.pt")
fused_model.eval()
fused_model.to(torch.device("cpu:0"))

img = cv2.imread("./1.png")
img_fp32 = img.astype(np.float32)
img_fp32 = (img_fp32-127.5) / 127.5
input = torch.from_numpy(img).permute(2, 0, 1).unsqueeze(0).float()

def speed_test(model, input):
    # warm up
    for i in range(10):
        model(input)

    import time
    start = time.time()
    for i in range(100):
        model(input)
    end = time.time()
    print("model time: ", (end-start)/100)
    time.sleep(10)

# quantized model
quantized_model= torch.jit.load("quantized_model.pt")
quantized_model.eval()
quantized_model.to(torch.device("cpu:0"))

speed_test(fused_model, input)
speed_test(quantized_model, input)

实测fp32模型单核运行120ms, 量化后47ms

结语

本文介绍了fx graph模式下的Pytorch的PTSQ方法，并实测了一个模型，效果还比较不错。