📚 课程目录第3章量化感知训练（QAT）

3. 量化感知训练（QAT）：伪量化节点、直通估计器（STE）、QAT训练流程与技巧

聊完PTQ，咱们得聊聊QAT了。说实话，PTQ虽然快，但精度损失有时候真让人头疼。我做过一个项目，模型在GPU上跑得好好的，一量化到INT8，准确率直接掉了3个点。业务方不干了，说这没法上线。怎么办？上QAT。

QAT的全称是Quantization-Aware Training，翻译过来就是“量化感知训练”。它的核心思想很简单：在训练过程中模拟量化带来的误差，让模型自己去适应这种误差。说白了，就是让模型提前“感受”一下被量化后的痛苦，然后自己学会怎么在这种痛苦下还能表现好。

核心观点：QAT不是训练完再量化，而是在训练过程中就“假装”已经量化了。这样模型学到的参数天然就对量化误差有鲁棒性。

3.1 伪量化节点（Fake Quantization Nodes）

QAT里最核心的概念就是伪量化节点。你可能会问，什么叫“伪量化”？

嗯，是这样的。真正的量化是把浮点数变成整数，这个过程是不可逆的——你丢了精度就找不回来了。但在训练中，我们需要梯度反向传播，整数操作没法求导啊。所以QAT搞了个折中：在前向传播时模拟量化的效果，但在反向传播时仍然使用浮点梯度。

伪量化节点做的事情，说白了就是三步：

量化：把浮点数映射到整数范围（比如INT8的[-128, 127]）
反量化：把整数再映射回浮点数
输出：得到的是浮点数，但已经带上了量化误差

这个过程可以用一个公式表示：

# 伪量化操作
def fake_quantize(x, scale, zero_point, bit_width=8):
    # 量化
    x_int = torch.round(x / scale + zero_point)
    # 截断到整数范围
    x_int = torch.clamp(x_int, 0, 2**bit_width - 1)
    # 反量化
    x_fp = (x_int - zero_point) * scale
    return x_fp

你看，输入和输出都是浮点数，但中间走了一趟整数。这样梯度就能正常计算了。我在项目中经常把伪量化节点插在卷积层后面、激活函数前面，效果最好。

我的经验：伪量化节点不是随便插的。我个人习惯在权重和激活值上都加，但偏置一般不加——偏置的分布太窄，量化后误差反而大。

3.2 直通估计器（STE）

好，现在问题来了。伪量化操作里有个torch.round，这个函数的梯度是多少？

答案是：几乎处处为0。因为round函数是个阶梯函数，它的导数除了在跳变点，其他地方都是0。如果梯度为0，那网络就学不动了。

这时候STE（Straight-Through Estimator）就登场了。它的想法很粗暴：前向传播用round，反向传播假装round不存在。也就是说，梯度直接穿透round函数，原封不动地传过去。

用数学表达就是：

# STE 的核心思想
# 前向: y = round(x)
# 反向: ∂L/∂x = ∂L/∂y  （假装round的导数是1）

我刚开始接触STE的时候，觉得这玩意儿也太不严谨了。但后来发现，它确实work。为什么？因为量化误差相对于激活值本身来说是小量，用STE近似梯度虽然不精确，但方向大致是对的。训练过程中，模型会慢慢学会避开那些梯度不准确的地方。

注意：STE不是万能的。如果量化位宽太低（比如2bit、3bit），STE的近似误差会很大，模型可能不收敛。我曾经在4bit量化上吃过亏，后来加了梯度缩放才稳住。

3.3 QAT训练流程

QAT的训练流程，我总结下来分四步。每一步都有坑，咱们一个一个说。

3.3.1 第一步：预训练或加载预训练模型

QAT很少从零开始训练。一般都是先有一个浮点模型，不管是自己训练的还是在ImageNet上预训练的。这个浮点模型就是我们的起点。

为什么？因为QAT本质上是在微调，让模型适应量化。如果从零开始，模型连基本特征都还没学会，你让它去适应量化误差，那真是难为它了。

3.3.2 第二步：插入伪量化节点

在模型的关键位置插入伪量化节点。哪些是关键位置？

权重：每个卷积层或全连接层的权重
激活值：每个卷积层或全连接层的输出
输入：模型的输入数据（可选，看情况）

我一般用PyTorch的torch.quantization工具包，它提供了QuantStub和DeQuantStub，可以很方便地插入伪量化节点。

import torch
import torch.nn as nn
import torch.quantization as quant

class QATModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.quant = quant.QuantStub()  # 输入量化
        self.conv1 = nn.Conv2d(3, 64, 3)
        self.relu = nn.ReLU()
        self.conv2 = nn.Conv2d(64, 128, 3)
        self.dequant = quant.DeQuantStub()  # 输出反量化

    def forward(self, x):
        x = self.quant(x)
        x = self.conv1(x)
        x = self.relu(x)
        x = self.conv2(x)
        x = self.dequant(x)
        return x

3.3.3 第三步：微调训练

这一步和普通训练差不多，但有几个关键区别：

训练参数	普通训练	QAT微调
学习率	通常较大（如1e-3）	很小（如1e-5 ~ 1e-6）
训练轮数	几十到几百轮	很少（通常5~10轮）
数据量	全量训练集	全量或子集
优化器	SGD/Adam	推荐SGD，动量小一点

为什么学习率要这么小？因为QAT是在一个已经很优的解附近微调，步子大了容易跨过最优解。我见过有人用1e-4的学习率做QAT，结果精度比PTQ还差。

3.3.4 第四步：导出量化模型

训练完成后，把伪量化节点去掉，换成真正的量化操作。这一步在PyTorch里叫convert。转换后，模型就变成INT8的了，可以部署到推理引擎上。

# 转换伪量化模型为真正的量化模型
model.eval()
model = quant.convert(model, inplace=False)
# 现在 model 就是 INT8 模型了

避坑指南：我曾经在导出时忘记设置model.eval()，结果BN层的统计量没冻结，导出的模型推理结果全是错的。记住，导出前一定要eval。

3.4 QAT训练技巧

做QAT有些技巧，是我在项目中一点点试出来的。分享几个我觉得最有用的。

3.4.1 逐层学习率调整

不同层对量化的敏感度不一样。一般来说，浅层（靠近输入的层）对量化误差更敏感，深层相对不敏感。所以我会给浅层更小的学习率，深层稍微大一点。

# 逐层设置学习率
optimizer = torch.optim.SGD([
    {'params': model.conv1.parameters(), 'lr': 1e-6},
    {'params': model.conv2.parameters(), 'lr': 5e-6},
    {'params': model.fc.parameters(), 'lr': 1e-5},
], lr=1e-5, momentum=0.9)

3.4.2 冻结BN层

QAT微调时，BN层的统计量（running_mean和running_var）最好冻结。因为微调的数据量少，BN统计量更新不稳定，反而会引入噪声。我一般在前几轮就冻结BN，后面只更新权重。

3.4.3 使用余弦退火学习率

QAT的训练轮数少，用阶梯式学习率衰减容易跳过头。我推荐用余弦退火（Cosine Annealing），它让学习率平滑下降，最后收敛得更稳。

scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(
    optimizer, T_max=10, eta_min=1e-7
)

3.4.4 校准数据的选择

QAT虽然用全量数据更好，但有时候数据量太大，训练时间受不了。这时候可以选一个子集做校准。选什么样的子集？我建议选那些分布多样、覆盖各种场景的数据。不要只选简单的样本，那样校准出来的量化参数对难样本不友好。

总结一下：QAT的核心就是“模拟量化误差 + STE近似梯度 + 小学习率微调”。它比PTQ慢，但精度通常更好。如果你的模型对精度要求极高，或者PTQ后精度掉得太多，那就上QAT。

3.5 QAT知识体系图

下面这张图把QAT的核心流程和关键概念串起来了。你可以看到，从浮点模型出发，经过插入伪量化节点、STE梯度近似、微调训练，最后导出INT8模型。每一步都有对应的技巧和注意事项。

嗯，QAT的内容差不多就这些。说实话，QAT是量化技术里最需要动手实践的部分。理论看着简单，但真正调起来，你会发现各种奇怪的问题——梯度爆炸、不收敛、精度反而下降。别慌，这些都是正常的。多试几次，找到适合你模型的那组参数，就好了。