第3章：训练后量化（PTQ）——实战中的精度调优与验证

训练后量化，也就是PTQ，是我在实际项目中用得最多的模型压缩手段。为什么？因为它快。你不需要重新训练模型，拿一个训练好的浮点模型，跑一遍校准流程，就能得到一个INT8模型。说白了，就是“白嫖”推理加速。

但这里有个坑——精度掉得厉害怎么办？我遇到过不止一次，模型量化后精度从90%掉到60%，那叫一个崩溃。所以这一章，咱们就聊聊PTQ的完整流程、校准数据集怎么选，以及在TensorRT和ONNX Runtime里具体怎么落地。

3.1 PTQ流程详解

PTQ的流程其实不复杂，但每一步都有讲究。我习惯把它拆成四个阶段：

1. 准备浮点模型：你得先有一个训练好的FP32模型。注意，模型必须是推理状态，不是训练状态。BatchNorm层要融合，Dropout要关掉。
2. 插入量化节点：在模型的每个权重和激活值后面，插入伪量化节点（FakeQuantize）。这些节点负责记录数据的动态范围。
3. 校准（Calibration）：跑一批校准数据，让量化节点收集到真实的激活值分布。这一步决定了量化参数——scale和zero_point。
4. 生成量化模型：根据校准结果，把FP32的权重和激活值映射到INT8，生成最终的量化模型。

嗯，这里要注意：校准不是训练。校准过程中，权重是不更新的。你只是让模型“看”一遍数据，然后统计出每个张量的最大值、最小值，或者更复杂的分布信息。

3.2 校准数据集的选择

校准数据集，说白了就是用来“喂”模型的那几百张图片或几百条数据。我见过有人随便拿100张训练集图片做校准，结果量化后精度崩了。为什么？因为校准集没有覆盖到模型可能遇到的各种情况。

我个人习惯，校准数据集的选择要遵循三个原则：

• 代表性：校准集的数据分布，要和实际推理时的数据分布一致。比如你做的是猫狗分类，校准集里就不能全是猫，没有狗。
• 多样性：尽量覆盖不同的光照、角度、背景。我曾在项目中只用白天拍的图片做校准，结果模型在夜间场景下直接“瞎了”。
• 适量性：不是越多越好。通常500-1000张图片就足够了。太多反而会让校准过程变慢，而且收益递减。

我曾经踩过一个坑：用训练集的全部数据做校准，结果量化后精度反而比用500张子集还差。后来分析发现，训练集里有很多重复或相似的样本，导致校准统计出来的动态范围偏大，量化精度反而下降。

3.3 PTQ在TensorRT中的实现

TensorRT的PTQ，我用的最多的是它的INT8校准器（Int8Calibrator）。TensorRT提供了几种校准算法：

校准算法	特点	适用场景
EntropyCalibrator	基于KL散度，寻找最优阈值	大多数分类模型，精度损失小
MinMaxCalibrator	直接取最大值作为阈值	对精度要求极高，但可能浪费量化范围
PercentileCalibrator	取某个百分位数的值	存在异常值的情况

我个人习惯，先用EntropyCalibrator跑一遍。如果精度掉得厉害，再换成PercentileCalibrator，调一下百分位数试试。

下面是一个简单的TensorRT PTQ代码示例：

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

# 创建builder和config
builder = trt.Builder(trt.Logger(trt.Logger.WARNING))
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_data, batch_size):
        trt.IInt8EntropyCalibrator2.__init__(self)
        self.calibration_data = calibration_data
        self.batch_size = batch_size
        self.current_index = 0
        # 分配GPU内存
        self.device_input = cuda.mem_alloc(self.calibration_data[0].nbytes)

    def get_batch(self, names):
        if self.current_index + self.batch_size > len(self.calibration_data):
            return None
        batch = self.calibration_data[self.current_index:self.current_index + self.batch_size]
        self.current_index += self.batch_size
        cuda.memcpy_htod(self.device_input, batch)
        return [int(self.device_input)]

    def get_batch_size(self):
        return self.batch_size

    def read_calibration_cache(self):
        return None

    def write_calibration_cache(self, cache):
        pass

# 使用校准器
calibrator = MyCalibrator(calibration_images, batch_size=32)
config.int8_calibrator = calibrator

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

嗯，这里要注意：校准器的get_batch方法，每次返回一个batch的数据。数据必须是连续的，而且类型要对。我遇到过有人传了float32的数据，结果校准出来的scale全是错的。

3.4 PTQ在ONNX Runtime中的实现

ONNX Runtime的PTQ，走的是另一种路子。它用的是ONNX的量化工具——onnxruntime.quantization。这个工具的好处是，你不需要写太多代码，几行就能搞定。

但坏处是，灵活性不如TensorRT。你没法自定义校准算法，只能用它内置的那几种。

下面是一个ONNX Runtime PTQ的示例：

import onnx
from onnxruntime.quantization import quantize_static, QuantType, CalibrationMethod
from onnxruntime.quantization import CalibrationDataReader

# 自定义校准数据读取器
class MyDataReader(CalibrationDataReader):
    def __init__(self, calibration_images):
        self.calibration_images = calibration_images
        self.current_index = 0

    def get_next(self):
        if self.current_index >= len(self.calibration_images):
            return None
        input_name = "input"
        data = self.calibration_images[self.current_index]
        self.current_index += 1
        return {input_name: data}

    def rewind(self):
        self.current_index = 0

# 执行静态量化
quantize_static(
    model_input="model.onnx",
    model_output="model_quant.onnx",
    calibration_data_reader=MyDataReader(calibration_images),
    quant_format=QuantType.QInt8,
    per_channel=True,
    activation_type=QuantType.QInt8,
    weight_type=QuantType.QInt8,
    calibrate_method=CalibrationMethod.MinMax,
    extra_options = {
        "ActivationSymmetric": True,
        "WeightSymmetric": True,
    }
)

你想想看，这个流程是不是比TensorRT简单多了？但简单也有简单的代价。ONNX Runtime的校准方法只有MinMax和Entropy两种，而且不支持自定义校准器。如果你的模型对量化特别敏感，可能得在TensorRT上多花点功夫。

我曾经在一个语义分割模型上，用ONNX Runtime的PTQ，精度掉了5个点。后来换成TensorRT的EntropyCalibrator，只掉了1个点。所以，选哪个工具，得看你的精度要求。

3.5 精度验证与调优

量化完了，怎么验证精度？很简单，跑一遍验证集，对比FP32和INT8的精度。如果精度掉得在可接受范围内（比如1-2%），那就OK。如果掉得太多，就得调优了。

我常用的调优手段有三个：

• 换校准算法：从Entropy换成Percentile，或者反过来。
• 调整校准集：增加或减少样本，或者换一批数据。
• 混合精度：把某些敏感层保留为FP16或FP32。在TensorRT里，可以通过设置per_layer_precision来实现。

嗯，最后一条最有效，但也最麻烦。你得一层一层地试，找到哪些层对量化最敏感。我习惯的做法是，先全部量化，然后逐层恢复为FP16，看精度变化。如果恢复某一层后精度明显回升，那这层就是“敏感层”。

好了，这一章就聊到这儿。下一章咱们聊聊量化感知训练（QAT），那才是真正能“救”模型精度的手段。