3、量化感知训练（QAT）：原理、伪量化节点与训练技巧

聊完PTQ，咱们得聊聊QAT了。说实话，PTQ虽然方便，但精度损失有时候真让人头疼。我去年有个项目，模型在FP32上跑得挺好，一量化到INT8，关键指标直接掉了3个点——这谁能忍？

这时候，QAT就派上用场了。它的核心思想很简单：让模型在训练过程中就「适应」量化带来的误差。说白了，就是提前告诉模型：「嘿，你以后要活在INT8的世界里，现在先练练手」。

3.1 QAT的原理：模拟量化误差

QAT的原理，用一句话概括就是：在前向传播时模拟量化，在反向传播时保持全精度梯度。

为什么会这样设计？你想想看，量化本身是个不可导的操作——round函数在绝大多数地方的导数都是0。如果直接对量化后的权重求导，梯度根本传不回去，模型就没法训练了。

所以QAT的做法是：

• 前向传播：把权重和激活值量化到INT8，再反量化回FP32，用这个「模拟量化」后的值做计算
• 反向传播：梯度直接跳过量化节点，用全精度的原始权重去更新

这个技巧叫做「直通估计器」（STE，Straight-Through Estimator）。我刚开始接触时觉得这方法有点「粗暴」，但实践证明它非常有效。你可以理解为：前向走的是「量化后的路」，反向却用「全精度的地图」来修正方向。

3.2 如何插入伪量化节点

伪量化节点（FakeQuantize）是QAT的核心组件。它不是一个真正的量化器，而是一个「模拟器」——把浮点数先量化再反量化，让数值带上量化误差。

在PyTorch中，插入伪量化节点通常有两种方式：

方式一：手动插入（适合有经验的玩家）

我个人习惯在关键层前后手动插入torch.quantization.FakeQuantize。这样控制力最强，但代码量也最大。

import torch
from torch.quantization import FakeQuantize

# 创建一个伪量化节点，scale=0.1, zero_point=0
fake_quant = FakeQuantize(
    observer=torch.quantization.MinMaxObserver,
    quant_min=-128, quant_max=127,
    dtype=torch.qint8,
    qscheme=torch.per_tensor_symmetric
)

# 在卷积层前插入
x = fake_quant(x)  # 模拟量化输入
x = conv(x)        # 全精度卷积计算
x = fake_quant(x)  # 模拟量化输出

方式二：使用QAT API（推荐）

PyTorch提供了torch.quantization.prepare_qat，它会自动在模型中插入伪量化节点。你只需要指定要量化的模块类型即可。

import torch.quantization as quant

# 定义量化配置
qconfig = quant.QConfig(
    activation=quant.MinMaxObserver.with_args(dtype=torch.quint8, qscheme=torch.per_tensor_affine),
    weight=quant.MinMaxObserver.with_args(dtype=torch.qint8, qscheme=torch.per_tensor_symmetric)
)

# 将模型转换为QAT模式
model.qconfig = qconfig
model = quant.prepare_qat(model, inplace=True)

# 开始训练
for epoch in range(num_epochs):
    for data, target in dataloader:
        output = model(data)
        loss = criterion(output, target)
        loss.backward()
        optimizer.step()

3.3 QAT的训练技巧

QAT训练和普通训练不太一样。我总结了几条实战经验，希望能帮你少走弯路。

技巧1：从预训练模型开始，不要从头训

QAT的目的是微调，不是从头学习。直接用预训练好的FP32模型做初始化，通常只需要训练几个epoch就能收敛。我曾经试过从头训QAT模型，训了50个epoch精度还是上不去——白白浪费了算力。

技巧2：使用较小的学习率

QAT的学习率一般要比普通训练小一个数量级。我常用的配置是：

训练阶段	学习率	优化器
普通训练（FP32）	1e-3	Adam
QAT微调	1e-4 ~ 5e-5	SGD（带动量）

为什么用SGD而不是Adam？我个人经验是，SGD的更新更「平滑」，不容易破坏量化尺度。当然，这不是绝对的，你可以根据实际情况调整。

技巧3：冻结BN层，或者使用较小的BN动量

QAT训练时，BatchNorm的统计量（running_mean和running_var）会随着伪量化节点的插入而抖动。我建议：

• 如果训练数据量足够大，可以冻结BN层（model.eval()模式）
• 如果数据量小，把BN的动量设小一些，比如momentum=0.01

技巧4：逐步放宽量化范围

一开始就把量化范围设得很窄（比如min=-1, max=1），模型可能学不好。我建议采用渐进式量化：

1. 先用较宽的量化范围（比如min=-10, max=10）训练几个epoch
2. 然后根据实际数据分布，逐步收紧范围
3. 最后固定量化参数，做最终微调

这个方法在MobileNet这类轻量级模型上特别有效。我记得有一次，用渐进式量化把精度从92.1%提到了92.8%，几乎追平了FP32的水平。

技巧5：注意量化后的算子融合

QAT训练完成后，别忘了做算子融合。常见的融合模式有：

• Conv + BN + ReLU → ConvReLU
• Conv + BN + ReLU6 → ConvReLU6
• Linear + ReLU → LinearReLU

PyTorch提供了torch.quantization.fuse_modules，可以自动完成融合。这一步能减少内存访问次数，提升推理速度。

# 算子融合示例
model = quant.fuse_modules(model, [['conv1', 'bn1', 'relu1']])
model = quant.convert(model, inplace=True)  # 转换为INT8

3.4 QAT vs PTQ：什么时候用哪个？

很多同学问我：「到底该用QAT还是PTQ？」我的建议是：

场景	推荐方法	理由
模型精度损失 < 0.5%	PTQ	省时省力，不需要重新训练
精度损失 > 1%	QAT	必须用QAT来弥补精度差距
模型部署到边缘设备	QAT	边缘设备对精度更敏感，QAT更稳妥
快速迭代验证	PTQ	先看PTQ效果，不行再上QAT

嗯，QAT的内容差不多就这些了。总结一下：QAT的核心是STE + 伪量化节点，训练时注意学习率、BN层和量化范围的设置。如果你在项目中遇到精度损失问题，不妨试试QAT——虽然多花点训练时间，但效果往往值得。