3. 量化感知训练(QAT):让模型学会“低精度生存”
大家好,我是你们的老朋友。今天我们来聊聊量化里最硬核的一块——量化感知训练(QAT)。
说实话,很多人一开始接触量化,都是从PTQ(训练后量化)入门的。毕竟简单,模型训完直接转,省事。但PTQ有个问题——精度掉得厉害,尤其是小模型或者对精度敏感的任务。这时候,QAT就派上用场了。
QAT的核心思想很简单:让模型在训练阶段就“体验”量化带来的误差,然后通过微调去适应它。说白了,就是让模型学会在低精度下“生存”。
3.1 QAT的原理:模拟量化,反向传播
QAT的原理,我可以用一句话概括:在前向传播时模拟量化,在反向传播时假装没量化。
为什么会这样?因为量化本身是个离散操作,比如把浮点数四舍五入到整数,这个操作不可导。如果直接对量化后的模型做反向传播,梯度传不回去。所以QAT用了一个技巧——直通估计器(STE,Straight-Through Estimator)。
STE的做法是:前向传播时,我们真的把权重和激活值量化到低精度(比如INT8),计算损失;反向传播时,我们假装量化操作不存在,直接把梯度传回给浮点权重。这样,浮点权重就能根据量化后的误差更新了。
嗯,这里要注意:更新的是浮点权重,不是量化后的整数权重。量化后的整数权重只是前向传播时“临时”用的。
我在项目中遇到过一个问题:如果学习率设置太大,浮点权重更新幅度过大,会导致量化后的整数权重频繁跳变,模型反而训不收敛。所以QAT的学习率通常要比普通训练小一个数量级。
3.2 伪量化节点(FakeQuant):QAT的核心组件
QAT的实现,靠的是在计算图中插入伪量化节点(FakeQuant)。这个节点干了两件事:
- 量化:把浮点数映射到低精度整数(比如INT8)
- 反量化:再把整数映射回浮点数(但精度已经损失了)
所以,伪量化节点的输出仍然是浮点数,但数值已经“模拟”了量化后的精度损失。模型看到的是浮点数,但感受到的是量化误差。
伪量化节点的数学表达是这样的:
# 伪量化操作
def fake_quantize(x, scale, zero_point, num_bits=8):
# 量化
x_int = torch.round(x / scale + zero_point)
# 截断到有效范围
x_int = torch.clamp(x_int, 0, 2**num_bits - 1)
# 反量化
x_fp = (x_int - zero_point) * scale
return x_fp
你看,输入是浮点,输出也是浮点,但中间经过了量化-反量化的过程。这就是“伪量化”名字的由来——看起来没变,实际上精度已经掉了。
我个人习惯在模型的权重和激活值后面都插入FakeQuant节点。权重量化通常用per-tensor(整个张量共享一个scale),激活值量化用per-tensor或者per-channel,看具体硬件支持。
关键点:FakeQuant节点只在训练时生效,推理时会被移除,替换成真正的量化算子。
3.3 QAT的完整流程:从浮点到INT8
QAT的流程,我总结为四步走:
- 准备预训练模型:先拿一个浮点模型,在原始任务上训好。这一步不能省,QAT不是从零开始训练。
- 插入FakeQuant节点:在模型的权重和激活值后面插入伪量化节点。注意,BN层(Batch Normalization)的统计量需要冻结,否则QAT过程中BN的均值和方差会乱掉。
- 微调训练:用较小的学习率(比如1e-5)继续训练几个epoch。数据量不需要太大,一般一个epoch就够了,但要看任务复杂度。
- 导出量化模型:训练完成后,把FakeQuant节点移除,导出真正的INT8模型(权重是INT8,激活值也是INT8)。
我曾经踩过一个坑:在第二步插入FakeQuant节点时,没有冻结BN层。结果训练过程中BN的running_mean和running_var不断更新,导致量化后的模型精度比PTQ还差。后来我学乖了,插入FakeQuant后第一件事就是冻结BN。
下面是一个简化的QAT训练代码示例:
import torch
import torch.nn as nn
import torch.quantization as quant
# 1. 准备预训练模型
model = MyModel()
model.load_state_dict(torch.load('pretrained.pth'))
# 2. 插入FakeQuant节点
model.qconfig = quant.get_default_qat_qconfig('fbgemm')
quant.prepare_qat(model, inplace=True)
# 冻结BN层
for module in model.modules():
if isinstance(module, nn.BatchNorm2d):
module.eval() # 冻结BN统计量
# 3. 微调训练
optimizer = torch.optim.SGD(model.parameters(), lr=1e-5)
for epoch in range(5):
for data, target in dataloader:
output = model(data)
loss = criterion(output, target)
loss.backward()
optimizer.step()
optimizer.zero_grad()
# 4. 导出量化模型
model.eval()
quantized_model = quant.convert(model, inplace=False)
torch.save(quantized_model.state_dict(), 'qat_model.pth')
注意,quant.convert这一步会把FakeQuant节点移除,替换成真正的量化算子。转换后的模型,权重就是INT8了。
3.4 QAT vs PTQ:什么时候该用谁?
这个问题,我经常被问到。我的建议是:能PTQ就别QAT,除非PTQ精度不够。
为什么?因为QAT需要训练,需要数据,需要调参,成本高。PTQ只需要一个校准集,跑一遍前向传播就完事了。
但PTQ有个硬伤:它对小模型和敏感任务不友好。比如MobileNet、BERT这类模型,PTQ后精度可能掉2-3个点,甚至更多。这时候QAT就香了。
我整理了一个对比表格,方便你决策:
| 对比维度 | PTQ(训练后量化) | QAT(量化感知训练) |
|---|---|---|
| 是否需要训练 | 不需要 | 需要微调 |
| 是否需要数据 | 少量校准集(几百张图) | 训练集(几千张图以上) |
| 精度损失 | 大模型0.5-1%,小模型2-5% | 通常0.1-0.5% |
| 部署成本 | 低,一步到位 | 高,需要训练+转换 |
| 适用场景 | 大模型、精度不敏感任务 | 小模型、精度敏感任务 |
我个人经验是:如果模型参数量超过100M,PTQ通常够用;如果模型小于10M,建议直接上QAT。中间那档,先试PTQ,精度不行再换QAT。
避坑指南:我曾经在一个OCR识别模型上试PTQ,精度从98%掉到92%,直接没法用。后来换成QAT,微调了3个epoch,精度回到97.5%。所以,别在PTQ一棵树上吊死。
3.5 QAT的进阶技巧:如何训得更好?
QAT虽然效果好,但也不是无脑训就能成的。我分享几个实战技巧:
- 学习率衰减:QAT的学习率建议从1e-5开始,每2个epoch衰减0.5倍。训太快容易震荡。
- 逐层量化:如果模型很大,可以逐层插入FakeQuant节点,先量化前几层,再量化后几层。这样梯度更稳定。
- 混合精度:有些层对量化特别敏感(比如第一层卷积),可以保留FP16或FP32,其他层用INT8。这叫混合精度量化。
- 校准集选择:QAT用的数据最好和训练集分布一致。如果用了分布不同的数据,模型会“学歪”。
我记得有一次,一个语音识别模型QAT后精度反而下降了。排查了半天,发现是校准集里混了噪声数据。换了干净数据后,精度立刻回升。所以,数据质量比数量更重要。
3.6 本章小结:QAT的核心要点
好了,我们来捋一捋QAT的核心要点:
- 原理:用STE解决量化不可导问题,让模型在训练时适应量化误差
- 组件:FakeQuant节点,模拟量化-反量化过程
- 流程:预训练 → 插入FakeQuant → 微调 → 导出INT8模型
- 对比:QAT精度高但成本高,PTQ成本低但精度可能不够
QAT不是银弹,但它是目前精度最高的量化方案。如果你的模型对精度要求苛刻,或者PTQ后精度掉得没法看,那就试试QAT吧。
最后,送大家一句话:量化不是目的,部署才是。选对方案,少走弯路。
警告:QAT训练完成后,一定要用quant.convert导出真正的量化模型。不要直接把带FakeQuant节点的模型拿去部署,那玩意儿推理速度比浮点还慢。
想要更多资料,联系可可