一、训练后量化（PTQ）：原理、校准与实战

大家好，我是你们的老朋友。今天咱们聊聊训练后量化，也就是 PTQ。

说实话，PTQ 是我在实际项目里用得最多的量化方式。为什么？因为它简单、直接、见效快。你不需要重新训练模型，不需要动你的训练脚本，甚至不需要原始数据集——只需要一小部分校准数据，就能把 FP32 模型压到 INT8，推理速度翻倍，内存砍半。

嗯，听起来很美好对吧？但坑也不少。我踩过的坑，今天一并告诉你。

1.1 PTQ 的核心原理

PTQ 的全称是 Post-Training Quantization，翻译过来就是「训练后量化」。说白了，就是模型训练完了，我再动手脚。

它的核心思路很简单：

  • 把 FP32 的权重和激活值，映射到 INT8 的整数范围
  • 这个映射需要找到一个「缩放因子」（scale）和「零点偏移」（zero_point）
  • 推理时，用整数运算代替浮点运算

你可能会问：为什么不直接训练 INT8 模型？因为训练时梯度需要高精度，INT8 根本扛不住。所以 PTQ 成了最务实的方案。

核心公式：

Q = round( R / scale ) + zero_point

其中 R 是浮点值，Q 是量化后的整数值，scale 和 zero_point 就是我们要找的参数。

我在项目中遇到过一个问题：模型量化后精度掉得厉害，从 98% 掉到 85%。后来发现是 scale 选得太粗糙了。所以，校准这一步，真的很关键。

1.2 校准数据集：选对数据，事半功倍

校准数据集，说白了就是用来「试探」模型激活值分布的一组数据。它不需要很大，几百到几千张图片就够了。

但这里有个坑：校准数据集必须和训练数据的分布一致。我曾经用 ImageNet 的验证集做校准，结果模型在工业场景下直接崩了——因为工业图片和自然图片的分布差太远了。

避坑指南：

我曾经用 100 张图片做校准，结果精度掉了 5%。后来换成 500 张，精度只掉了 0.5%。所以，校准集的数量不能太少，建议至少 200-500 张。

校准数据集的选择原则：

  • 覆盖所有类别
  • 包含边缘情况（比如光照变化、遮挡等）
  • 数量适中，不要太多（否则校准时间太长）
  • 不要用训练集（容易过拟合到训练分布）

1.3 KL 散度与熵校准：找到最优的 scale

校准的核心，就是找到每个张量的 scale 和 zero_point。最常用的方法之一，就是基于 KL 散度的校准。

KL 散度，全称 Kullback-Leibler divergence，用来衡量两个分布之间的差异。在量化中，我们用它来衡量原始浮点分布和量化后分布之间的差异。

具体做法是这样的：

  1. 收集激活值的直方图（比如 2048 个 bin）
  2. 尝试不同的阈值（threshold），把直方图压缩到 256 个 bin
  3. 计算原始分布和压缩后分布的 KL 散度
  4. 选择 KL 散度最小的那个阈值

你想想看，这个过程其实就是在找一个「最优的截断点」。把那些离群值截断掉，让量化后的分布尽可能接近原始分布。

个人经验：

我习惯用 KL 散度校准激活值，用 min-max 校准权重。因为权重的分布通常比较集中，min-max 就够了；而激活值的分布变化很大，需要 KL 散度来精细调整。

为什么 KL 散度效果好？因为它对分布的尾部很敏感。那些出现频率低但数值很大的离群值，如果直接保留，会拉大 scale，导致大部分值量化精度下降。KL 散度会自动「砍掉」这些离群值，换来整体精度的提升。

1.4 PTQ 实战：PyTorch 和 TensorFlow

好了，理论说完了，咱们直接上代码。

PyTorch 实现 PTQ

PyTorch 从 1.8 版本开始，提供了非常方便的量化工具。我个人习惯用 torch.quantization 模块。

import torch
import torch.quantization as quant

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

# 配置量化后端
model.qconfig = quant.get_default_qconfig('fbgemm')

# 准备校准
model = quant.prepare(model, inplace=True)

# 校准：用校准数据集跑一遍 forward
for images, _ in calib_loader:
    model(images)

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

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

嗯，就这么简单。但要注意几点：

  • 校准前一定要 model.eval()，否则 BN 层会出问题
  • 校准数据不要用梯度，用 torch.no_grad() 包一下
  • fbgemm 是 x86 平台的量化后端，arm 平台用 qnnpack

TensorFlow 实现 PTQ

TensorFlow 这边，我用的是 TFLite 的量化工具。

import tensorflow as tf

# 加载模型
converter = tf.lite.TFLiteConverter.from_saved_model('model_savedmodel')

# 配置量化
converter.optimizations = [tf.lite.Optimize.DEFAULT]
converter.representative_dataset = representative_dataset

# 转换为量化模型
tflite_model = converter.convert()

# 保存
with open('model_int8.tflite', 'wb') as f:
    f.write(tflite_model)

这里的 representative_dataset 就是校准数据集生成器。它需要返回一个迭代器，每次 yield 一个输入样本。

def representative_dataset():
    for images, _ in calib_loader.take(200):
        yield [images.numpy()]

关键点：

TensorFlow 的 PTQ 默认使用 KL 散度校准。如果你想用 min-max，可以手动设置 tf.lite.RepresentativeDataset 的量化参数。

1.5 知识体系图

下面这张图，是我自己总结的 PTQ 知识体系。你可以把它当作一个快速索引。

训练后量化（PTQ）知识体系 PTQ 原理 校准数据集 KL 散度与熵校准 FP32 → INT8 映射 scale & zero_point 整数运算加速 200-500 张样本 覆盖所有类别 与训练分布一致 直方图收集 阈值搜索 分布差异最小化 PyTorch 实战 TensorFlow 实战 INT8 量化模型

1.6 常见问题与避坑

最后，分享几个我踩过的坑：

  • 校准数据太少： 我曾经用 50 张图片做校准，结果精度掉了 8%。后来加到 500 张，精度只掉了 0.3%。
  • 忘记 eval 模式： 校准前没设 model.eval()，BN 层还在更新，结果量化后模型直接崩了。
  • 后端选错： x86 平台用了 qnnpack，推理速度反而慢了。记得用 fbgemm。
  • 校准数据分布不一致： 用自然图片校准，部署到工业场景，精度掉得厉害。一定要用目标场景的数据。

我的建议：

PTQ 是性价比最高的量化方式。如果你的模型精度掉得不多（<1%），直接用 PTQ 就够了。如果掉得厉害，再考虑 QAT（量化感知训练）。

好了，PTQ 的内容就到这里。记住一句话：校准数据集选得好，量化成功一半。剩下的，就是调调 scale 和 zero_point 了。

可可精品资料库

← 上一章 📖 返回目录 下一章 →