4、QAT训练流程：准备预训练模型、插入伪量化节点、微调训练、导出量化模型
好，咱们进入正题。量化感知训练，也就是QAT，说白了就是让模型在训练阶段就“学会”忍受量化带来的精度损失。我刚开始接触这个的时候，总觉得多此一举——直接训练完再量化不就行了？后来被现实狠狠教育了一回。
嗯，今天咱们就把QAT的完整流程拆开揉碎了讲。一共四步：准备预训练模型、插入伪量化节点、微调训练、导出量化模型。每一步都有坑，我踩过的坑尤其多。
4.1 准备预训练模型
这一步其实最容易被忽视。很多人上来就拿一个随便训好的模型开始做QAT，结果效果惨不忍睹。我个人习惯是，先确保预训练模型的精度已经足够好。你想想看，如果原模型精度只有80%，量化后掉到75%，那还能接受。但如果原模型精度本来就只有70%，量化后掉到60%，那就没法用了。
我在项目中遇到过一件事：有个同事拿了一个还没收敛的模型直接做QAT，折腾了两周，精度死活上不去。后来我让他先把模型训到收敛，再跑QAT，一天就搞定了。所以，预训练模型的质量，决定了QAT的天花板。
核心要点：预训练模型必须已经充分收敛，且精度达到业务要求。不要指望QAT能帮你把没训好的模型救回来。
具体来说，准备预训练模型时，我建议你检查以下几点：
- 模型是否已收敛？看验证集loss是否还在下降，如果还在下降，说明没训完。
- 精度是否达标？如果原模型精度就不满足需求，QAT也救不了。
- 模型结构是否支持量化？有些自定义算子可能不支持量化，需要提前替换。
小技巧：我一般会在预训练阶段就保存好checkpoint，包括优化器状态和learning rate schedule。这样后面做QAT微调时，可以直接从最佳checkpoint继续，省去重新加载的麻烦。
4.2 插入伪量化节点
这一步是QAT的核心。伪量化节点，英文叫FakeQuantize，说白了就是在模型里插入一些“假量化”操作。它模拟了量化-反量化的过程，让模型在训练时就能感受到量化带来的误差。
为什么会这样？因为直接量化模型时，权重和激活值从浮点数变成整数，精度损失是突然发生的。而QAT通过伪量化节点，把这个损失分散到训练过程中，让模型慢慢适应。
伪量化节点一般插在哪里？我总结了一下：
| 插入位置 | 说明 | 常见情况 |
|---|---|---|
| 权重之后 | 对卷积层、全连接层的权重做伪量化 | 几乎所有量化方案都需要 |
| 激活值之后 | 对ReLU、sigmoid等激活函数的输出做伪量化 | 大部分量化方案需要 |
| 残差连接处 | 对残差相加的结果做伪量化 | ResNet、Transformer等结构需要 |
| concat操作处 | 对拼接后的张量做伪量化 | DenseNet、YOLO等结构需要 |
在PyTorch中，插入伪量化节点其实很简单。官方提供了torch.quantization.QuantStub和torch.quantization.DeQuantStub，配合torch.quantization.prepare_qat函数，可以自动完成插入。代码示例如下：
import torch
import torch.nn as nn
import torch.quantization
# 定义模型
class MyModel(nn.Module):
def __init__(self):
super().__init__()
self.quant = torch.quantization.QuantStub()
self.conv = nn.Conv2d(3, 64, 3)
self.relu = nn.ReLU()
self.dequant = torch.quantization.DeQuantStub()
def forward(self, x):
x = self.quant(x) # 输入量化
x = self.conv(x)
x = self.relu(x)
x = self.dequant(x) # 输出反量化
return x
# 准备QAT
model = MyModel()
model.qconfig = torch.quantization.get_default_qat_qconfig('fbgemm')
model = torch.quantization.prepare_qat(model, inplace=True)注意：我曾经犯过一个错误——在插入伪量化节点后，忘记设置qconfig。结果模型跑起来跟没量化一样，白白浪费了两天训练时间。所以，一定要检查model.qconfig是否被正确设置。
4.3 微调训练
伪量化节点插好了，接下来就是微调训练。这一步跟普通训练很像，但有几点不同。
第一，学习率要调小。我一般会把学习率降到原来的十分之一到百分之一。因为模型已经收敛了，太大的学习率会破坏已有的权重分布。我记得有一次，我用原来的学习率直接跑QAT，结果loss直接炸了，模型精度掉到比随机猜还低。
第二，训练轮数不用太多。通常几个epoch就够了。我个人的经验是，对于ImageNet级别的数据集，5-10个epoch就能看到效果。对于小数据集，2-3个epoch可能就够了。
第三，要监控量化相关的指标。除了常规的loss和精度，我还会关注伪量化节点的scale和zero_point。如果这些值在训练过程中剧烈变化，说明模型还没稳定，需要继续训练。
微调训练的代码示例：
# 微调训练
optimizer = torch.optim.SGD(model.parameters(), lr=1e-4) # 学习率调小
criterion = nn.CrossEntropyLoss()
for epoch in range(5): # 少量epoch
for batch in dataloader:
inputs, labels = batch
optimizer.zero_grad()
outputs = model(inputs)
loss = criterion(outputs, labels)
loss.backward()
optimizer.step()
# 打印量化参数，观察是否稳定
for name, module in model.named_modules():
if hasattr(module, 'scale'):
print(f'{name} scale: {module.scale.item():.4f}')避坑指南：我曾经在微调时用了weight decay，结果量化后的精度反而下降了。后来发现，weight decay会改变权重的分布，导致量化参数不稳定。所以，我建议在QAT微调阶段，把weight decay设小一点，甚至设为0。
4.4 导出量化模型
微调完成后，最后一步就是导出量化模型。这一步其实包含两个子步骤：
- 融合算子：把BN层和卷积层融合在一起，减少计算量。
- 转换模型：把伪量化节点变成真正的量化操作，生成量化后的模型。
在PyTorch中，这两步可以通过torch.quantization.convert函数完成：
# 融合算子
model.eval()
model = torch.quantization.fuse_modules(model, [['conv', 'relu']])
# 转换模型
model = torch.quantization.convert(model, inplace=True)
# 保存量化模型
torch.save(model.state_dict(), 'quantized_model.pth')导出后的模型，权重和激活值都是int8类型。你可以用ONNX导出，也可以直接用PyTorch的JIT trace。我个人习惯导出ONNX，因为部署时更灵活。
重要提醒：导出量化模型后，一定要做精度验证。我见过太多人导出后直接部署，结果精度掉了5个点都没发现。验证时，用跟预训练模型一样的测试集，对比精度差异。如果差异在1%以内，基本可以接受。
嗯，到这里，QAT的完整流程就讲完了。总结一下：
- 准备预训练模型：确保模型已收敛，精度达标。
- 插入伪量化节点：在权重和激活值之后插入FakeQuantize，模拟量化误差。
- 微调训练：用小学习率、少epoch训练，让模型适应量化。
- 导出量化模型：融合算子、转换模型、验证精度。
每一步都有细节，每一步都可能踩坑。但只要你按照这个流程走，QAT其实没那么难。下次咱们聊聊更进阶的话题——如何在QAT中处理自定义算子。到时候我再分享几个我踩过的坑，保证让你少走弯路。