3、量化感知训练（QAT）原理：模拟量化操作、直通估计器（STE）、训练策略与技巧
各位同学，咱们今天聊聊量化感知训练，也就是 QAT。说实话，这是我从 PTQ 转到 QAT 时，感觉最「玄学」的一块。PTQ 像搭积木，数据喂进去，模型出来，简单粗暴。但 QAT 不一样，它更像是在训练过程中「植入」了一个量化的大脑，让模型自己学会适应低精度。
我刚开始接触 QAT 时，最大的困惑是：训练时明明是浮点，怎么模拟量化？梯度怎么传回去？ 别急，咱们一个一个拆开看。
3.1 模拟量化操作：前向传播的「假动作」
QAT 的核心思想，就是在训练的前向传播中，插入一个模拟量化操作。说白了，就是假装把浮点权重和激活值量化成 INT8，再反量化回浮点。这样模型在训练时就能「感受」到量化带来的精度损失，并主动调整参数去补偿。
这个操作通常用 torch.quantization.FakeQuantize 实现。它的数学表达很简单：
# 模拟量化：浮点 -> 量化 -> 反量化
def fake_quantize(x, scale, zero_point, qmin=0, qmax=255):
# 1. 量化
x_int = torch.round(x / scale + zero_point)
x_int = torch.clamp(x_int, qmin, qmax)
# 2. 反量化
x_fp = (x_int - zero_point) * scale
return x_fp
嗯，这里要注意：模拟量化操作本身不改变数值类型，输入输出都是 FP32。它只是在前向时引入了量化误差，让模型「提前适应」部署时的精度损失。
关键点：模拟量化操作只在训练时生效。推理时，我们会把 FakeQuantize 节点移除，替换成真正的量化算子（比如 INT8 卷积）。
我在项目中遇到过一个问题：模拟量化操作放在哪里？ 通常放在权重和激活值之后、卷积/全连接层之前。但具体位置会影响精度，我建议你多试几种配置，别死磕一种。
3.2 直通估计器（STE）：让梯度「绕道而行」
好，现在问题来了。模拟量化操作中有一个 torch.round，这个函数在反向传播时梯度是 0（因为 round 的导数几乎处处为 0）。梯度为 0，参数就更新不了，训练就废了。
怎么办？直通估计器（STE） 就是来解决这个问题的。它的思想很简单：前向传播用 round，反向传播假装 round 不存在，直接把梯度传过去。
用代码实现就是：
class FakeQuantizeSTE(torch.autograd.Function):
@staticmethod
def forward(ctx, x, scale, zero_point):
# 前向：正常模拟量化
x_int = torch.round(x / scale + zero_point)
x_int = torch.clamp(x_int, 0, 255)
x_fp = (x_int - zero_point) * scale
return x_fp
@staticmethod
def backward(ctx, grad_output):
# 反向：梯度直通，不做任何修改
return grad_output, None, None
说白了，STE 就是一个「骗」梯度的技巧。它让梯度绕过了 round 这个不可微的操作，直接更新底层的浮点权重。你想想看，如果没有 STE，QAT 根本跑不起来。
个人经验：STE 虽然简单，但有个隐含假设——量化误差足够小，梯度近似正确。如果量化误差太大（比如从 FP32 直接量化到 INT4），STE 的效果会变差。我曾经在 4-bit 量化时遇到过梯度爆炸，后来加了梯度裁剪才稳住。
3.3 训练策略与技巧：QAT 不是简单微调
QAT 的训练策略，和普通训练有很大区别。我总结了几条实战经验，供你参考。
3.3.1 初始化：从 PTQ 开始
我个人习惯，先用 PTQ 做一次校准，得到 scale 和 zero_point 的初始值。然后把这些参数加载到 QAT 模型中，再开始训练。这样能避免 QAT 初期因为量化参数不准确导致的剧烈震荡。
具体做法：
- 用 PTQ 校准 100-200 个 batch，得到每层的 scale 和 zero_point。
- 把这些参数冻结，作为 QAT 的初始值。
- 开始 QAT 训练，同时微调 scale 和 zero_point（通常用指数移动平均更新）。
3.3.2 学习率：小一点，再小一点
QAT 的学习率通常比普通训练小 10-100 倍。我一般用 1e-5 到 1e-4 起步。为什么？因为 QAT 是在一个已经收敛的模型上做「微调」，学习率太大容易破坏原有特征。
我曾经试过用 1e-3 的学习率跑 QAT，结果验证集精度直接掉了 5 个点。后来改成 1e-5，慢慢回升了。
3.3.3 训练时长：别太长，也别太短
QAT 的训练 epoch 数，通常是原始训练的 10%-20%。比如原始训练 100 个 epoch，QAT 跑 10-20 个 epoch 就够了。跑太久反而会过拟合到量化噪声上。
我建议你每 2-3 个 epoch 做一次验证，如果精度不再提升，就提前停止。
3.3.4 量化参数的更新策略
scale 和 zero_point 的更新，有两种常见策略：
| 策略 | 做法 | 适用场景 |
|---|---|---|
| 冻结 | 训练过程中不更新 scale/zero_point | 量化误差小，模型稳定 |
| EMA 更新 | 用指数移动平均更新 scale/zero_point | 量化误差大，需要动态调整 |
我个人更推荐 EMA 更新，尤其是做 8-bit 以下量化时。EMA 的 momentum 通常设为 0.99 或 0.999。
3.4 避坑指南：我踩过的那些坑
最后，分享几个我实际踩过的坑，希望能帮你少走弯路。
坑 1： 我曾经在 QAT 训练时，忘记把 BN 层设置为 eval 模式。结果 BN 层的 running_mean 和 running_var 被不断更新，导致推理时精度崩了。记住：QAT 中 BN 层通常保持冻结，用预训练好的统计量。
坑 2： 模拟量化操作中的 scale 和 zero_point，如果初始化不合理，会导致训练初期 loss 剧烈震荡。我建议先用 PTQ 校准 100 个 batch，再开始 QAT。
坑 3： 有些框架（比如 PyTorch）的 FakeQuantize 默认使用 per-tensor 量化。如果你部署时用的是 per-channel 量化，记得在 QAT 中也改成 per-channel，否则精度会不一致。
好了，QAT 的原理和策略就讲到这里。下一章咱们聊聊量化部署的实战技巧，包括如何用 TensorRT 和 ONNX Runtime 跑量化模型。到时候我会分享一些调优经验，保证让你少踩坑。