使用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方法，并实测了一个模型，效果还比较不错。
