pytorch FX模型静态量化

前言

以前面文章写到的mobilenet图像分类为例，本文主要记录一下pytorchh训练后静态量化的过程。

一、pytorch静态量化（手动版）

静态量化是最常用的量化形式，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量化

1.版本

fx量化版本也有坑，之前在torch 1.7版本操作总是报错搞不定，换成1.12.0版本就正常了，这一点非常重要。

2.代码如下：

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

总结

刚开始搞量化坑比较多，一个是某些操作不支持，合并层麻烦，另外有版本问题导致的报错可能搞很久，觉得有用的各位吴彦祖麻烦送个免费三连