以前面文章写到的mobilenet图像分类为例，本文主要记录一下pytorchh训练后静态量化的过程。
静态量化是最常用的量化形式，float32的模型量化成int8，模型大小大概变为原来的1/4，推理速度我在intel 8700k CPU上测试速度正好快4倍，但是在AMD的5800h CPU 上测试速度反而慢了两倍，应该是AMD不支持某些指令集加速。
之前手动添加量化节点的方式搞了好几天，最后模型是出来了，但是推理时候报错，大多数时候是RuntimeError: Could not run ‘--------’ with arguments from the ‘CPU’ backend，网上是说推理的时候没有安插QuantStub()和DeQuantStub()，可能是用的这个MobilenetV3网络结构复杂，某些地方没有手动添加到，这种方式肯定是可以成功的，只是比较麻烦容易出错。
# 加载模型
model = MobileNetV3_Large(2).to(device) # 加载一个网络，我这边是二分类传了一个2
checkpoint = torch.load(weights, map_location=device)
model.load_state_dict(checkpoint)
model.to('cpu').eval()
合并层对于一些重复使用的Block和nn.Sequential要打印出来看，然后append到mix_list 里面
比如
# 打印model
for name, module in model.named_children():
print(name, module)
比如这里Sequential里面存在conv+bn+relu，append进去的应该是[‘bneck.0.conv1’, ‘bneck.0.bn1’,‘nolinear1’]，但是nolinear1是个变量，也就是说某些时候是relu某些时候又不是，这种时候就要一个个分析判断好然后写代码，稍微复杂点就容易出错或者遗漏。
backend = "fbgemm" # x86平台
model.qconfig = torch.quantization.get_default_qconfig(backend)
mix_list = [['conv1','bn1'], ['conv2','bn2']] # 合并层只支持conv+bn conv+relu conv+bn+relu等操作，具体可以查一下,网络中存在的这些操作都append到mix_list里面
model = torch.quantization.fuse_modules(model,listmix) # 合并某些层
model_fp32_prepared = torch.quantization.prepare(model)
model_int8 = torch.quantization.convert(model_fp32_prepared)
有时候存在不支持的操作relu6这些要替换成relu，加法操作也要替换，最后还要输入一批图像校准模型等
self.skip_add = nn.quantized.FloatFunctional()
# forward的时候比如return a+b 改为return self.skip_add.add(a, b)
一系列注意事项操作完毕，最后推理各种报错，放弃了
fx量化版本也有坑，之前在torch 1.7版本操作总是报错搞不定，换成1.12.0版本就正常了，这一点非常重要。
import torch
import torch.nn as nn
import torch.nn.functional as F
import copy
import torchvision
from torchvision import transforms
from torch.quantization.quantize_fx import prepare_fx, convert_fx
from torch.quantization import get_default_qconfig
from torch import optim
import os
import time
from utils import load_data
from models.mobilenetv3copy import MobileNetV3_Large
def evaluate_model(model, test_loader, device, criterion=None):
model.eval()
model.to(device)
running_loss = 0
running_corrects = 0
for inputs, labels in test_loader:
inputs = inputs.to(device)
labels = labels.to(device)
outputs = model(inputs)
_, preds = torch.max(outputs, 1)
if criterion is not None:
loss = criterion(outputs, labels).item()
else:
loss = 0
# statistics
running_loss += loss * inputs.size(0)
running_corrects += torch.sum(preds == labels.data)
eval_loss = running_loss / len(test_loader.dataset)
eval_accuracy = running_corrects / len(test_loader.dataset)
return eval_loss, eval_accuracy
def quant_fx(model, data_loader):
model_to_quantize = copy.deepcopy(model)
model_to_quantize.eval()
qconfig = get_default_qconfig("fbgemm")
qconfig_dict = {"": qconfig}
prepared_model = prepare_fx(model_to_quantize, qconfig_dict)
print("开始校准")
calibrate(prepared_model, data_loader) # 这是输入一批有代表性的数据来校准
print("校准完毕")
quantized_model = convert_fx(prepared_model) # 转换
return quantized_model
def calibrate(model, data_loader):
model.eval()
with torch.no_grad():
for image, target in train_loader:
model(image)
if __name__ == "__main__":
cuda_device = torch.device("cuda:0")
cpu_device = torch.device("cpu:0")
model = MobileNetV3_Large(2) # 加载自己的网络
train_loader, test_loader = load_data(64, 8) # 自己写一个pytorch加载数据的方法
# quantization
state_dict = torch.load('./mymodel.pth') # 加载一个正常训练好的模型
model.load_state_dict(state_dict)
model.to('cpu')
model.eval()
quant_model = quant_fx(model, train_loader) # 执行量化代码
quant_model.eval()
print("开始验证")
eval_loss, eval_accuracy = evaluate_model(model=quant_model,
test_loader=test_loader,
device=cpu_device,
criterion=nn.CrossEntropyLoss())
print("Epoch: {:02d} Eval Loss: {:.3f} Eval Acc: {:.3f}".format(
-1, eval_loss, eval_accuracy))
torch.jit.save(torch.jit.script(quant_model), 'outQuant.pth') # 保存量化后的模型
# 加载量化模型推理
loaded_quantized_model = torch.jit.load('outQuant.pth')
eval_loss, eval_accuracy = evaluate_model(model=quant_model,
test_loader=test_loader,
device=cpu_device,
criterion=nn.CrossEntropyLoss())
print("Epoch: {:02d} Eval Loss: {:.3f} Eval Acc: {:.3f}".format(
-1, eval_loss, eval_accuracy))
fx量化也不用管里面什么算子不支持之类的，开箱即用，以上代码参考pytorch官网https://pytorch.org/docs/stable/fx.html
最后验证模型精度下降0.02%可以忽略不计，pytorch量化的模型是不支持gpu推理的，只能在arm或者x86平台实现压缩提速。要用cuda的话要上tensorrt+onnx，以后完成了再讲。完整的训练模型量化模型的代码后面会放到github上面。
完整代码:https://github.com/Ysnower/pytorch-static-quant
刚开始搞量化坑比较多，一个是某些操作不支持，合并层麻烦，另外有版本问题导致的报错可能搞很久，觉得有用的各位吴彦祖麻烦送个免费三连