4. 训练后量化（PTQ）实战：PTQ流程，校准数据集的选择，使用TensorRT/PyTorch进行PTQ

训练后量化，简称PTQ，是嵌入式部署里最常用的量化手段。说白了，就是模型训练完了，我们直接拿它做量化，不需要重新训练。我刚开始接触量化时，总觉得PTQ是“偷懒”的做法，后来才发现，它其实是工程落地的最佳平衡点——精度损失小，操作又简单。

4.1 PTQ流程：三步走，不迷路

PTQ的流程，我习惯把它拆成三步。你想想看，就像做菜一样，先备料、再烹饪、最后调味。量化也是这个道理。

1. 第一步：准备校准数据集——这是“备料”。选几组有代表性的输入数据，让模型“看一眼”真实数据的分布。
2. 第二步：执行量化——这是“烹饪”。用校准数据跑一遍前向推理，同时统计每一层的激活值范围（min/max），然后计算缩放因子和零点。
3. 第三步：验证与调优——这是“调味”。量化后的模型跑一下验证集，看看精度掉了多少。如果掉得太多，就得调整校准策略或者换数据集。

嗯，这里要注意：第二步里，统计激活值范围时，不同的校准方法（比如MinMax、Entropy、Percentile）会得到不同的结果。我个人习惯先用MinMax快速试水，如果精度不行，再换Entropy。

4.2 校准数据集的选择：成败在此一举

校准数据集，说白了就是给模型“打个样”。你拿什么数据给它看，它就按什么数据来量化。所以，选错数据集，量化效果直接崩盘。

我曾经踩过一个坑：用ImageNet的验证集做校准，结果模型在工业场景下精度掉了5个点。后来才发现，校准集的分布和实际部署场景的分布不一致。从那以后，我坚持一个原则：校准数据集必须来自部署场景的真实数据。

校准数据集类型	适用场景	注意事项
训练集子集	通用场景，数据分布均匀	避免选太多同类样本，要覆盖所有类别
验证集	快速验证，精度要求不高	可能过拟合到验证集分布
真实部署数据	工业场景，精度要求高	需要提前采集，数据量不用太大

为什么会这样？因为量化时，模型会统计每一层激活值的范围。如果校准集里全是“晴天”的图片，模型就不知道“阴天”的数据该怎么量化。结果就是，阴天图片的激活值被截断，精度直接崩了。

4.3 使用TensorRT进行PTQ

TensorRT的PTQ，我用的最多。它提供了两种方式：一种是直接调用API，另一种是用trtexec命令行工具。我个人习惯用API，因为可以灵活控制校准过程。

下面是一个完整的TensorRT PTQ示例。注意，这里用的是Python API，C++ API也类似。

import tensorrt as trt
import pycuda.driver as cuda
import pycuda.autoinit
import numpy as np

# 创建builder和config
TRT_LOGGER = trt.Logger(trt.Logger.WARNING)
builder = trt.Builder(TRT_LOGGER)
network = builder.create_network(1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH))
config = builder.create_builder_config()

# 设置量化模式为INT8
config.set_flag(trt.BuilderFlag.INT8)

# 设置校准器
class MyCalibrator(trt.IInt8EntropyCalibrator2):
    def __init__(self, calibration_files, batch_size, input_shape):
        trt.IInt8EntropyCalibrator2.__init__(self)
        self.calibration_files = calibration_files
        self.batch_size = batch_size
        self.input_shape = input_shape
        self.batch_idx = 0
        self.device_input = cuda.mem_alloc(batch_size * np.prod(input_shape) * 4)

    def get_batch_size(self):
        return self.batch_size

    def get_batch(self, names):
        if self.batch_idx < len(self.calibration_files):
            # 读取校准数据
            batch = []
            for i in range(self.batch_size):
                img = np.load(self.calibration_files[self.batch_idx + i])
                batch.append(img)
            batch = np.array(batch, dtype=np.float32)
            cuda.memcpy_htod(self.device_input, batch)
            self.batch_idx += self.batch_size
            return [int(self.device_input)]
        else:
            return None

    def read_calibration_cache(self, length):
        return None

    def write_calibration_cache(self, cache):
        return

# 创建校准器实例
calibrator = MyCalibrator(calibration_files, batch_size=32, input_shape=(3, 224, 224))
config.int8_calibrator = calibrator

# 解析ONNX模型
parser = trt.OnnxParser(network, TRT_LOGGER)
with open("model.onnx", "rb") as f:
    parser.parse(f.read())

# 构建引擎
engine = builder.build_engine(network, config)

# 保存引擎
with open("model_int8.engine", "wb") as f:
    f.write(engine.serialize())

这段代码里，最关键的是MyCalibrator类。它负责提供校准数据，并统计激活值范围。我习惯用EntropyCalibrator2，因为它对精度损失的控制更好。如果模型对精度特别敏感，可以试试MinMaxCalibrator，虽然简单，但有时效果反而更好。

4.4 使用PyTorch进行PTQ

PyTorch的PTQ，主要靠torch.quantization模块。它提供了两种模式：Eager Mode和FX Graph Mode。我个人更推荐FX Graph Mode，因为它自动化程度高，不容易出错。

下面是一个PyTorch PTQ的完整示例。

import torch
import torch.quantization as quant

# 加载预训练模型
model = torch.load("model.pth")
model.eval()

# 准备量化配置
model.qconfig = quant.get_default_qconfig('fbgemm')
# 或者用更激进的配置
# model.qconfig = quant.get_default_qconfig('qnnpack')

# 准备校准数据集
calibration_loader = torch.utils.data.DataLoader(
    calibration_dataset,
    batch_size=32,
    shuffle=True
)

# 执行PTQ
model = quant.prepare(model, inplace=False)

# 用校准数据跑一遍前向推理
with torch.no_grad():
    for images, _ in calibration_loader:
        model(images)

# 转换为量化模型
model = quant.convert(model, inplace=False)

# 保存量化模型
torch.save(model.state_dict(), "model_quantized.pth")

这段代码里，prepare函数负责插入量化/反量化节点，convert函数负责把权重和激活值转为INT8。注意，校准数据跑一遍前向推理时，模型会统计每一层的激活值范围。这个步骤不能少，否则量化后的模型精度会崩。

4.5 TensorRT vs PyTorch PTQ：怎么选？

你可能会问，TensorRT和PyTorch的PTQ，到底选哪个？我的建议是：看部署平台。

• 如果部署在NVIDIA GPU上，无脑选TensorRT。它的INT8 kernel优化得更好，推理速度更快。
• 如果部署在CPU或ARM上，选PyTorch的PTQ。它支持fbgemm和qnnpack后端，对CPU友好。
• 如果模型结构复杂（比如有自定义算子），PyTorch的FX Graph Mode可能解析不了，这时可以试试TensorRT的ONNX路径。

我个人习惯是：先用PyTorch PTQ快速验证量化效果，如果精度达标，再转TensorRT做最终部署。这样既能快速迭代，又能保证最终性能。