3、PTQ实战：使用PyTorch进行Post-Training Quantization

模型量化，说白了就是把模型里那些32位浮点数（FP32）的权重和激活值，用更少的位数来表示。最常见的就是INT8量化。为什么要这么做？因为模型变小了，推理变快了，功耗也降低了。在工业缺陷检测这种需要实时响应的场景里，这几乎是必选项。

我个人习惯把量化方法分成两类：训练后量化（PTQ）和量化感知训练（QAT）。这一章我们只聊PTQ，因为它最省事——不需要重新训练模型，拿一个训练好的模型，直接做量化校准就行。

3.1 PTQ的核心流程

PTQ的流程其实不复杂。我把它拆成四个步骤，你照着做就行。

1. 准备校准数据集——从训练集里抽几百张图，不需要标签，只需要覆盖常见的缺陷类型和背景。
2. 插入量化/反量化节点——在模型的关键位置插入Observer，用来统计激活值的分布范围。
3. 运行校准——把校准数据跑一遍，Observer会记录每个激活张量的min/max值。
4. 转换模型——根据统计到的范围，计算scale和zero_point，把FP32模型转成INT8模型。

嗯，这里要注意：校准数据不能太少，我建议至少200张。我在一个项目中试过只用50张，结果量化后的模型在某个缺陷类别上直接掉了5个点。后来换成500张，精度损失才降到0.3%。

3.2 PyTorch中的量化工具

PyTorch从1.3版本开始就内置了量化工具。我主要用这几个模块：

模块/函数	作用	我的使用建议
torch.quantization.QuantStub	在模型入口插入量化节点	放在第一个卷积层之前
torch.quantization.DeQuantStub	在模型出口插入反量化节点	放在最后一个全连接层之后
torch.quantization.FakeQuantize	模拟量化效果，用于校准	一般不需要手动调用，由prepare自动插入
torch.quantization.prepare	插入Observer，准备校准	调用前确保模型已经调用了.eval()
torch.quantization.convert	将校准后的模型转为INT8	转换后模型不可再训练

你想想看，这些工具其实帮我们做了大部分脏活。我们只需要定义好模型结构，然后调用几个API就行。

3.3 实战：量化一个缺陷检测模型

下面我拿一个简单的分类模型来演示。假设我们有一个训练好的ResNet18，用来检测PCB板上的缺陷。

import torch
import torch.nn as nn
import torch.quantization as quant

# 1. 定义模型，插入量化节点
class QuantizedResNet(nn.Module):
    def __init__(self, model_fp32):
        super().__init__()
        self.quant = quant.QuantStub()
        self.model = model_fp32
        self.dequant = quant.DeQuantStub()
    
    def forward(self, x):
        x = self.quant(x)
        x = self.model(x)
        x = self.dequant(x)
        return x

# 2. 加载预训练模型
model_fp32 = torch.load('pcb_defect_resnet18.pth')
model = QuantizedResNet(model_fp32)
model.eval()

# 3. 配置量化方案
model.qconfig = quant.get_default_qconfig('fbgemm')
# 或者用 per_channel 量化，精度更高
# model.qconfig = quant.get_default_qconfig('fbgemm_per_channel')

# 4. 准备校准
model_prepared = quant.prepare(model, inplace=False)

# 5. 运行校准
def calibrate(model, calib_loader):
    model.eval()
    with torch.no_grad():
        for images, _ in calib_loader:
            model(images)

calibrate(model_prepared, calib_loader)

# 6. 转换为INT8
model_int8 = quant.convert(model_prepared, inplace=False)

# 7. 保存量化模型
torch.save(model_int8.state_dict(), 'pcb_defect_resnet18_int8.pth')

3.4 量化配置的选择

PyTorch提供了几种默认的量化配置。我整理了一个对比表：

配置名称	量化粒度	适用后端	精度损失	我的推荐场景
fbgemm	per tensor	x86 CPU	0.5%~1.5%	通用场景，兼容性好
fbgemm_per_channel	per channel	x86 CPU	0.2%~0.8%	对精度敏感的场景
qnnpack	per tensor	ARM CPU	0.5%~1.5%	移动端/嵌入式部署

我曾经在一个项目中，用fbgemm配置量化后，模型在某个缺陷类别上的召回率从92%掉到了87%。换成fbgemm_per_channel后，召回率恢复到了91%。所以，如果硬件支持，我建议优先用per_channel量化。

3.5 量化后的模型结构

下面这张图展示了量化前后模型结构的变化。你可以看到，FP32模型里的Conv2d和Linear被替换成了量化版本，中间还插入了量化/反量化节点。

3.6 精度评估与调优

量化完成后，第一件事就是跑一遍验证集，看看精度变化。我一般会对比三个指标：

• mAP（平均精度均值）——对于检测模型，这是最关键的指标。
• 各类别AP——看看有没有某个缺陷类别掉得特别厉害。
• 推理延迟——量化后应该至少快2倍以上。

如果发现精度损失超过1%，我会尝试以下调优方法：

1. 增加校准数据量——从500张增加到1000张，覆盖更多边缘情况。
2. 改用per_channel量化——对权重按通道量化，精度更高。
3. 跳过敏感层——有些层对量化特别敏感，可以保留为FP32。
4. 使用QAT——如果PTQ实在搞不定，就只能上量化感知训练了。

3.7 量化模型的部署

量化后的模型可以直接用PyTorch的JIT导出，然后部署到生产环境。我常用的导出方式：

# 导出为TorchScript
scripted_model = torch.jit.script(model_int8)
scripted_model.save('pcb_defect_int8.pt')

# 或者导出为ONNX
dummy_input = torch.randn(1, 3, 224, 224)
torch.onnx.export(
    model_int8,
    dummy_input,
    'pcb_defect_int8.onnx',
    opset_version=13,
    input_names=['input'],
    output_names=['output']
)

导出后，可以用ONNX Runtime或者TensorRT来加速推理。我实测过，在Intel Xeon CPU上，INT8模型比FP32模型快2.5倍左右，而精度只掉了0.4%。

3.8 本章小结

PTQ是工业缺陷检测部署中最实用的量化方案。它不需要重新训练，只需要少量校准数据，就能把模型压缩到原来的1/4，推理速度提升2~3倍。PyTorch提供了完整的工具链，从插入量化节点到校准再到转换，一套流程走下来，半小时就能搞定。