4. 量化参数实战：基于PyTorch的静态量化实现、校准数据集的选择

静态量化，说白了就是先把模型跑一遍，收集好激活值的分布信息，然后再做量化。跟动态量化那种“边跑边算”的方式不一样，静态量化把计算量大的部分提前做好了。我个人习惯把静态量化叫做“一次校准，终身受益”。

嗯，这里要注意。静态量化不是简单的把权重和激活值都转成int8就完事了。你得先理解它的核心流程：

1. 准备一个校准数据集
2. 把模型插入量化/反量化节点
3. 跑几遍校准数据，收集激活值的统计信息
4. 根据统计信息计算量化参数（scale和zero_point）
5. 冻结量化参数，导出量化模型

我在项目中遇到过不少同学，上来就直接调用torch.quantization.quantize，结果模型精度掉得一塌糊涂。为什么？说白了就是校准数据集没选对。

4.1 PyTorch静态量化三步走

PyTorch的静态量化，我把它拆成三个步骤。你跟着走一遍，基本不会出大问题。

4.1.1 第一步：准备模型

先把你的模型定义好，然后调用torch.quantization.prepare。这个函数会在模型里插入QuantStub和DeQuantStub节点，还会把nn.ReLU、nn.Conv2d这些算子替换成可观测的版本。

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

class MyModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.quant = quant.QuantStub()
        self.conv1 = nn.Conv2d(3, 64, 3)
        self.relu1 = nn.ReLU()
        self.conv2 = nn.Conv2d(64, 128, 3)
        self.relu2 = nn.ReLU()
        self.fc = nn.Linear(128 * 6 * 6, 10)
        self.dequant = quant.DeQuantStub()

    def forward(self, x):
        x = self.quant(x)
        x = self.conv1(x)
        x = self.relu1(x)
        x = self.conv2(x)
        x = self.relu2(x)
        x = x.view(x.size(0), -1)
        x = self.fc(x)
        x = self.dequant(x)
        return x

model = MyModel()
model.eval()
model.qconfig = quant.get_default_qconfig('fbgemm')
model_prepared = quant.prepare(model)

这里有个坑。我曾经在项目里忘了调用model.eval()，结果prepare之后模型还在训练模式，BatchNorm的行为完全不对。校准出来的量化参数全是错的。你想想看，校准数据跑一遍，激活值分布都变了，后面量化能准吗？

4.1.2 第二步：校准

校准这一步，说白了就是让模型“见见世面”。你把校准数据集喂进去，模型会记录每一层激活值的min和max（或者直方图分布）。

def calibrate(model, data_loader, num_batches=100):
    model.eval()
    with torch.no_grad():
        for i, (inputs, _) in enumerate(data_loader):
            if i >= num_batches:
                break
            model(inputs)

校准的时候，我建议你跑100到200个batch。太少的话，统计信息不够稳定；太多的话，边际收益递减。我在项目里试过跑500个batch，结果跟100个batch的精度几乎一样，白白浪费了时间。

4.1.3 第三步：转换

校准完之后，调用torch.quantization.convert。这个函数会把模型里的浮点算子替换成量化算子，同时把之前收集到的量化参数冻结到模型里。

model_quantized = quant.convert(model_prepared)

转换完之后，你就可以用torch.jit.script或者torch.jit.trace导出量化模型了。我个人习惯用trace，因为它对动态图的处理更稳定。

4.2 校准数据集的选择——成败的关键

校准数据集选不好，量化精度直接崩。这不是危言耸听。我在项目里见过有人拿ImageNet的验证集做校准，结果模型在真实场景下精度掉了5个点。为什么？因为验证集的数据分布跟真实场景不一样。

具体来说，有几点要注意：

• 数据来源：最好从真实业务数据里随机抽取。不要用训练集，训练集的数据分布往往跟推理场景有偏差。
• 数据量：500到2000张图片（或样本）就够了。太多反而不好，因为校准过程会记录所有激活值的极值，样本太多容易引入异常值。
• 类别平衡：如果是个分类模型，校准数据集里每个类别都要有。我曾经犯过这个错，校准集里全是猫，结果模型在狗上精度掉了8个点。
• 数据增强：校准的时候不要做数据增强。你想想看，校准的目的是收集真实的激活值分布，不是让模型学泛化能力。加了随机裁剪、翻转，分布就变了。

校准数据集类型	适用场景	推荐数量	注意事项
真实业务数据	生产环境部署	500-2000张	必须覆盖所有常见场景
验证集子集	快速实验验证	100-500张	不要跟测试集重叠
合成数据	无真实数据可用	1000-5000张	分布可能不准确，谨慎使用

4.3 量化参数的计算方式

PyTorch支持两种量化参数计算方式：minmax和histogram。默认是minmax，但我在项目里发现，histogram往往效果更好。

• minmax：直接取激活值的全局最小值和最大值。简单粗暴，但容易受异常值影响。如果校准数据里有一张特别亮的图片，整个量化范围就被拉大了，精度会下降。
• histogram：统计激活值的直方图分布，然后通过KL散度找到最优的量化范围。这种方式能自动忽略掉那些极少出现的异常值。

你可以通过设置qconfig来切换：

# 使用histogram方式
qconfig = quant.QConfig(
    activation=quant.MinMaxObserver.with_args(dtype=torch.quint8, qscheme=torch.per_tensor_affine),
    weight=quant.MinMaxObserver.with_args(dtype=torch.qint8, qscheme=torch.per_tensor_symmetric)
)
model.qconfig = qconfig

4.4 静态量化知识体系

下面这张图，是我自己总结的静态量化核心流程。你把它记在脑子里，以后做量化就不会迷路。

4.5 避坑指南

做静态量化，有几个坑我踩过，你注意避开：

4.6 完整代码示例

最后，给你一个完整的静态量化代码模板。你直接复制过去，改改模型和数据加载器就能用。

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

def static_quantize(model, calibration_loader, num_batches=100):
    """
    静态量化完整流程
    Args:
        model: 待量化的模型
        calibration_loader: 校准数据加载器
        num_batches: 校准batch数
    Returns:
        量化后的模型
    """
    model.eval()
    model.qconfig = quant.get_default_qconfig('fbgemm')

    # 第一步：准备
    model_prepared = quant.prepare(model)

    # 第二步：校准
    with torch.no_grad():
        for i, (inputs, _) in enumerate(calibration_loader):
            if i >= num_batches:
                break
            model_prepared(inputs)

    # 第三步：转换
    model_quantized = quant.convert(model_prepared)

    return model_quantized

# 使用示例
# model = MyModel()
# quantized_model = static_quantize(model, calibration_loader, num_batches=100)
# torch.jit.save(torch.jit.script(quantized_model), 'quantized_model.pt')

嗯，静态量化这块内容就这些。你只要把校准数据集选好，三步走完，基本不会出大问题。记住，校准数据集的质量决定了量化精度的上限。别在这个环节偷懒。