4. 量化感知训练(QAT)
量化感知训练,圈里人常叫 QAT。说白了,就是让模型在训练阶段就「体验」一下量化后的精度损失,然后自己想办法把损失补回来。
我刚开始接触 QAT 时,总觉得这玩意儿有点玄学。明明推理时要做量化,为什么训练时还要模拟一遍?后来在部署一个 MobileNetV3 的项目里踩了坑,才真正理解它的价值。嗯,咱们今天就把这块硬骨头啃下来。
4.1 为什么需要 QAT?
先看一个简单的问题。你用 PTQ(训练后量化)把模型从 FP32 压到 INT8,结果精度掉了 3 个点。怎么办?
有人会说:「调一下校准集,换一种量化方法试试。」但很多时候,这些招数都不管用。尤其是小模型、低比特量化(4-bit、2-bit),PTQ 基本扛不住。
QAT 的思路完全不同——它让模型在训练过程中「学会」量化带来的误差,然后通过梯度更新把权重调整到更适合量化后的分布。
核心思想:QAT 不是对训练好的模型做量化,而是让模型在训练时就「知道」自己将来会被量化。
4.2 模拟量化:前向传播的「假动作」
QAT 的前向传播,其实是在做一件很巧妙的事——用浮点数模拟整数的计算过程。
具体来说,我们在网络中插入 伪量化节点(FakeQuant)。这些节点干两件事:
- 量化:把浮点数映射到整数范围(比如 INT8 的 [-128, 127])
- 反量化:再把整数映射回浮点数
你想想看,这一来一回,精度损失就模拟出来了。但关键是,数据格式仍然是 FP32,梯度可以正常传播。
伪量化的数学表达式很简单:
# 伪量化操作
q = clamp(round(x / scale) + zero_point, q_min, q_max)
x_hat = (q - zero_point) * scale
这里的 scale 和 zero_point 就是量化参数。在 QAT 训练过程中,它们可以是固定的(从校准集统计得到),也可以是可训练的(LSQ 方法)。
我的经验:对于 8-bit 量化,用固定 scale 通常就够了。但如果是 4-bit 或更低,建议用可训练的 scale,精度能再拉回来 1-2 个点。
4.3 直通估计器(STE):反向传播的「障眼法」
这里有个大问题。你仔细看上面的伪量化公式——round() 函数,它的梯度几乎处处为 0。如果梯度为 0,网络就学不动了。
怎么办?
学术界提出了一个非常巧妙的 trick,叫 直通估计器(Straight-Through Estimator, STE)。
STE 的核心思想就一句话:前向传播该量化就量化,反向传播假装量化不存在。
具体来说:
- 前向:计算
y = round(x) - 反向:计算
∂y/∂x = 1(假装 round 是恒等函数)
这听起来有点粗暴,但实践证明它 works。为什么?因为 round 的误差其实很小,直接用 1 近似梯度,网络依然能收敛到不错的位置。
注意:STE 不是万能的。当量化比特数太低(比如 2-bit)时,round 的误差太大,STE 的近似就不太准了。我曾经在 2-bit 量化上试过,模型直接不收敛……后来改用 LSQ + 逐层调学习率才搞定。
4.4 QAT 的完整训练流程
咱们把整个流程串起来,看看 QAT 到底怎么跑:
| 步骤 | 操作 | 说明 |
|---|---|---|
| 1 | 加载预训练 FP32 模型 | 最好用已经收敛的模型,不要从头训 |
| 2 | 插入 FakeQuant 节点 | 在权重和激活值后面各加一个 |
| 3 | 初始化量化参数 | 用少量校准数据统计 scale/zero_point |
| 4 | 小学习率微调 | 通常用 1e-5 ~ 1e-4,训几个 epoch |
| 5 | 冻结 BN 层 | BN 在 QAT 中容易不稳定,建议冻结 |
| 6 | 导出 INT8 模型 | 去掉 FakeQuant,只保留量化权重 |
关键参数:QAT 的学习率一般设为原始训练的 1/10 到 1/100。太大容易破坏预训练权重,太小又学不动。
4.5 QAT 与 PTQ 的对比
很多同学会问:「QAT 这么麻烦,我直接用 PTQ 不行吗?」
我列个表,你一看就明白了:
| 维度 | PTQ | QAT |
|---|---|---|
| 是否需要训练 | 不需要 | 需要微调 |
| 精度损失 | 较大(尤其小模型) | 较小(可接近 FP32) |
| 部署时间 | 几分钟 | 几小时到几天 |
| 适用场景 | 大模型、高比特 | 小模型、低比特 |
| 硬件适配 | 通用 | 需支持 QAT 算子 |
我个人习惯是:先跑 PTQ,精度达标就直接用。不达标再上 QAT。别一上来就 QAT,浪费时间。
4.6 QAT 的工程实现细节
纸上谈兵差不多了,咱们看看代码。下面是一个 PyTorch 风格的伪量化实现:
class FakeQuantize(torch.autograd.Function):
@staticmethod
def forward(ctx, x, scale, zero_point, q_min, q_max):
# 量化
x_int = torch.round(x / scale) + zero_point
x_int = torch.clamp(x_int, q_min, q_max)
# 反量化
x_hat = (x_int - zero_point) * scale
return x_hat
@staticmethod
def backward(ctx, grad_output):
# STE:梯度直接通过
return grad_output, None, None, None, None
你看,反向传播里直接返回 grad_output,这就是 STE 的精髓——假装 round 不存在。
避坑指南:我曾经在实现时忘了处理 zero_point 的梯度。如果你的 scale 是可训练的,记得把 zero_point 也设成可训练,否则精度会差一截。
4.7 QAT 的常见变体
QAT 不是只有一种玩法。这些年学术界搞出了不少变体,我挑几个实用的说说:
- LSQ (Learned Step Size Quantization):让 scale 也参与训练,精度更高,但训练更不稳定
- PACT (Parameterized Clipping Activation):在激活值量化前加一个可训练的截断阈值,适合 ReLU 类激活函数
- QAT with Knowledge Distillation:用 FP32 教师模型指导 INT8 学生模型训练,效果拔群
我在部署一个 3-bit 的语义分割模型时,就是用了 LSQ + 知识蒸馏,才把 mIoU 从 68% 拉回到 72%。单用 QAT 只能到 70%。
4.8 QAT 的注意事项
最后,说几个容易踩的坑:
- BN 层要冻结:QAT 微调时,BN 的统计量会漂移,建议在微调前先 freeze BN
- 学习率不能太大:我见过有人用 1e-3 训 QAT,结果权重直接飞了
- 校准集要覆盖分布:QAT 用的数据最好和部署场景一致,否则量化参数会偏
- 逐层验证精度:不要只看整体精度,有时候某一层的量化误差会拖垮全局
嗯,QAT 的内容差不多就这些。说白了,它就是用训练时间换推理精度。如果你的模型对延迟和功耗要求极高,QAT 几乎是必经之路。