4、训练后量化(PTQ)：PTQ原理、校准数据集、KL散度与熵校准、实战：使用TensorRT进行PTQ。

各位同学，欢迎来到第四章。今天咱们聊聊训练后量化，也就是PTQ。

说实话，在嵌入式平台上做模型部署，PTQ是我最常用的手段。为什么？因为它省事啊。你不需要重新训练模型，拿一个训练好的浮点模型，跑一遍校准，就能得到一个量化模型。我刚开始接触这个的时候，总觉得“哪有这么好的事”，后来发现，只要校准数据集选得好，精度损失真的可以控制得很小。

4.1 PTQ原理：说白了就是找映射关系

PTQ的核心思想，就是把浮点数映射到定点数。比如一个32位浮点数的权重，我们想把它变成8位整数。怎么变？

公式其实很简单：

q = round(r / scale) + zero_point

其中r是原始浮点值，q是量化后的整数值，scale是缩放因子，zero_point是零点偏移。

嗯，这里要注意。scale和zero_point怎么确定？这就引出了校准过程。我个人习惯把校准理解为“找一把合适的尺子”。这把尺子要能量尽所有的数据，但又不能浪费刻度。

举个例子。假设你的权重范围是[-1.0, 1.0]，那scale就是2.0/255≈0.0078。但如果你的权重范围是[-0.5, 0.5]，还用同样的scale，那有一半的整数刻度就浪费了。所以校准的目的，就是找到最适合当前数据分布的scale和zero_point。

4.2 校准数据集：别小看它，坑很多

校准数据集，说白了就是用来“教”量化算法如何选择scale和zero_point的数据。

我在项目中遇到过一个问题。有个同学拿训练集直接当校准集，结果量化后的模型在测试集上精度掉得一塌糊涂。为什么？因为训练集和测试集的分布有差异，校准出来的尺子只适合训练集。

我个人建议，校准数据集应该满足几个条件：

• 代表性：能覆盖模型在实际部署中可能遇到的各种情况
• 多样性：不要全是同一类样本，比如全是猫的图片
• 数量适中：一般500-2000张图片就够了，太多反而过拟合

我曾经踩过一个坑。用100张图片做校准，结果量化后模型在某些场景下完全失效。后来换成500张，问题就解决了。你想想看，100张图片能代表多少种情况？

4.3 KL散度与熵校准：量化界的“黄金搭档”

说到校准方法，KL散度绝对是绕不开的话题。很多框架，包括TensorRT，默认用的就是KL散度校准。

KL散度衡量的是两个分布之间的差异。在量化中，我们用它来比较原始浮点值的分布和量化后整数值的分布。我们希望这两个分布尽可能接近。

具体怎么做？

1. 统计原始浮点值的直方图
2. 尝试不同的阈值，把直方图截断
3. 对截断后的直方图进行量化
4. 计算量化后的分布和原始分布的KL散度
5. 选择KL散度最小的那个阈值

说白了，就是找一个“最佳截断点”。为什么要截断？因为8位整数只有256个刻度，而浮点数的范围可能很大。如果硬要把所有浮点数都映射到256个刻度上，那每个刻度代表的区间就很大，精度损失就大。截断掉一些极端值，反而能让大部分数据的量化更精细。

我记得有一次，一个模型量化后精度掉了5个点。我检查了一下，发现是某个层的权重分布有很长的尾巴。用KL散度校准后，精度只掉了0.5个点。这就是截断的好处。

4.4 实战：使用TensorRT进行PTQ

好了，理论讲完了，咱们来点实际的。TensorRT做PTQ，其实就几步。

首先，你需要一个训练好的浮点模型。我这里以ONNX格式为例。

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()

# 设置量化标志
config.set_flag(trt.BuilderFlag.INT8)

# 设置校准器
class MyCalibrator(trt.IInt8EntropyCalibrator2):
    def __init__(self, calibration_files, batch_size):
        trt.IInt8EntropyCalibrator2.__init__(self)
        self.calibration_files = calibration_files
        self.batch_size = batch_size
        self.batch_idx = 0
        self.device_input = cuda.mem_alloc(self.batch_size * 3 * 224 * 224 * 4)
        
    def get_batch_size(self):
        return self.batch_size
    
    def get_batch(self, names):
        if self.batch_idx < len(self.calibration_files):
            # 读取并预处理图片
            batch = self._load_batch()
            cuda.memcpy_htod(self.device_input, batch)
            self.batch_idx += 1
            return [int(self.device_input)]
        else:
            return None
    
    def read_calibration_cache(self):
        return None
    
    def write_calibration_cache(self, cache):
        with open('calibration.cache', 'wb') as f:
            f.write(cache)

# 创建校准器实例
calibrator = MyCalibrator(calibration_files, batch_size=32)
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)

这段代码看着长，其实核心就几行：

• config.set_flag(trt.BuilderFlag.INT8)：告诉TensorRT我们要做INT8量化
• config.int8_calibrator = calibrator：设置校准器
• builder.build_engine()：构建量化后的引擎

校准完成后，你会得到一个INT8精度的TensorRT引擎。这个引擎的大小大约是FP32模型的1/4，推理速度能提升2-4倍。

最后，验证一下量化后的精度。我习惯的做法是：

1. 用FP32模型跑一遍测试集，记录精度
2. 用INT8模型跑一遍同样的测试集，记录精度
3. 对比两个精度，如果差距在1%以内，就算合格

如果精度掉得太多，怎么办？别急，有几个排查方向：

• 检查校准数据集是否足够有代表性
• 尝试不同的校准算法（比如MinMax校准）
• 对某些敏感层进行回退（保留为FP16）

好了，这一章的内容就到这里。PTQ是个实用技术，但也不是万能的。下一章我们会聊聊量化感知训练（QAT），那才是真正能“榨干”模型性能的手段。