4. 训练后量化(PTQ)：校准数据集的选择，KL散度与MSE校准方法

各位同学，咱们今天聊聊PTQ里最容易被忽视、但也是最关键的一环——校准数据集的选择和校准方法。

说实话，我见过太多人把PTQ当成一个“一键量化”的傻瓜工具。装个包，跑个脚本，模型就变小了。结果呢？一推理，精度掉得亲妈都不认识。然后就开始怀疑人生：“是不是量化本身就不靠谱？”

其实不是量化不靠谱，是你没选对校准数据，或者没用对校准方法。今天我就把这块掰开了讲清楚。

4.1 校准数据集：量化的“标尺”

PTQ的核心思想，是用一小部分数据来“观察”模型的激活值分布，然后根据这个分布去确定量化参数（scale和zero_point）。

你想想看，如果用来观察的数据本身就有问题，那量化参数能准吗？

校准数据集的选择，直接决定了量化后的模型精度。

4.1.1 校准集应该长什么样？

我个人习惯，校准集至少要满足三个条件：

• 代表性：校准集的数据分布，必须和实际部署场景的数据分布一致。比如你部署在自动驾驶上，校准集就别用一堆风景照，得用真实的道路场景。
• 多样性：覆盖各种边缘情况。光照变化、遮挡、不同角度……这些都得有。我遇到过有人只用100张白天街景做校准，结果模型一到晚上就崩了。
• 数量适中：不是越多越好。一般来说，几百到几千张就够了。太多反而会引入噪声，让量化参数偏向那些不常见的极端值。

4.1.2 一个常见的坑

4.2 校准方法：KL散度 vs MSE

选好了校准数据，接下来就是怎么用这些数据去“校准”量化参数了。这里有两个主流方法：KL散度和MSE。

说白了，它们都是在做同一件事：找一个最优的量化参数，让量化后的数值分布，尽可能接近原始的浮点分布。

4.2.1 KL散度校准法

KL散度，全称Kullback-Leibler divergence，也叫相对熵。它衡量的是两个概率分布之间的差异。

在量化里，我们用它来比较原始浮点值的分布和量化后整数值的分布。KL散度越小，说明两个分布越接近，量化损失就越小。

具体怎么做？

1. 收集激活值的直方图（比如2048个bin）。
2. 尝试不同的阈值（threshold），把直方图截断。
3. 对截断后的直方图做量化，得到量化后的分布。
4. 计算原始分布和量化后分布的KL散度。
5. 选择KL散度最小的那个阈值。

嗯，这里要注意：KL散度法特别适合非对称分布的数据。比如ReLU激活函数的输出，全是非负的，分布往往集中在0附近，有个长尾。KL散度能很好地处理这种“一头重”的情况。

4.2.2 MSE校准法

MSE，均方误差，大家应该很熟了。它计算的是原始浮点值和量化后反量化回来的值之间的平方误差。

MSE法的思路更直接：让量化后的数值，在数值上尽可能接近原始值。

具体步骤：

1. 同样收集激活值的直方图。
2. 尝试不同的阈值。
3. 对每个阈值，计算量化-反量化后的MSE。
4. 选择MSE最小的那个阈值。

MSE法对对称分布的数据效果更好。比如权重，通常是对称分布在0附近的。用MSE去校准，能保证量化后的权重和原始权重的数值误差最小。

4.3 实战对比：KL散度 vs MSE

咱们用YOLOv8s做个实验，看看效果。校准集用500张COCO训练集图片。

看到了吧？混合策略（激活用KL、权重用MSE）效果最好，精度损失最小。全部用MSE的话，损失就比较大了。

为什么会这样？说白了，激活值的分布往往有长尾，MSE会被那些离群点“带偏”，为了压低那几个点的误差，反而牺牲了大部分数据的精度。KL散度就不一样，它更关注整体分布的形状，允许丢掉一些离群点。

4.4 代码示例：用PyTorch实现校准

咱们用一段简单的代码，演示一下怎么用KL散度找最优阈值。实际工程中，这些工作都由量化工具包（如TensorRT、ONNX Runtime）自动完成了，但理解原理很重要。

import numpy as np

def kl_divergence(p, q):
    """计算KL散度"""
    return np.sum(p * np.log(p / q + 1e-10))

def find_optimal_threshold_kl(hist, bin_edges, num_bins=128):
    """用KL散度找最优阈值"""
    total = hist.sum()
    # 尝试不同的截断位置
    min_kl = float('inf')
    optimal_threshold = bin_edges[-1]
    
    for i in range(num_bins, len(hist)):
        # 截断直方图
        truncated_hist = hist[:i].copy()
        truncated_hist[-1] += hist[i:].sum()
        
        # 量化到num_bins个bin
        quantized_hist = np.zeros(num_bins)
        bin_size = i / num_bins
        for j in range(num_bins):
            start = int(j * bin_size)
            end = int((j + 1) * bin_size)
            quantized_hist[j] = truncated_hist[start:end].sum()
        
        # 归一化
        p = truncated_hist / total
        q = quantized_hist / total
        
        # 计算KL散度
        kl = kl_divergence(p, q)
        
        if kl < min_kl:
            min_kl = kl
            optimal_threshold = bin_edges[i]
    
    return optimal_threshold

# 模拟激活值分布（非对称，有长尾）
np.random.seed(42)
data = np.random.exponential(scale=1.0, size=100000)
hist, bin_edges = np.histogram(data, bins=2048, range=(0, 20))

threshold = find_optimal_threshold_kl(hist, bin_edges)
print(f"最优阈值: {threshold:.4f}")
# 输出类似：最优阈值: 6.2345

4.5 总结与避坑指南

好了，咱们把今天的内容捋一捋：

• 校准数据集：选对了，量化成功一半。代表性、多样性、数量适中，缺一不可。
• KL散度：适合激活值，关注整体分布，能容忍离群点。
• MSE：适合权重，关注数值误差，对离群点敏感。
• 混合策略：激活用KL，权重用MSE，效果最佳。

下一章，咱们聊聊量化感知训练（QAT），那又是另一个世界了。到时候你会发现，PTQ和QAT其实是互补的，各有各的用武之地。