量化感知训练（QAT）原理：模拟量化、伪量化节点、直通估计器（STE）、QAT训练流程

大家好，欢迎来到第二章。上一章我们聊了PTQ（训练后量化），那种“一刀切”的方式虽然快，但精度损失有时候让人头疼。尤其是当模型要部署到边缘设备上，精度掉个1%都受不了。

这时候，就该QAT（量化感知训练）登场了。说白了，QAT就是在训练过程中，让模型“提前适应”量化带来的噪声。我个人的习惯是，如果PTQ掉点超过0.5%，我就会毫不犹豫地切换到QAT。今天，我们就来彻底拆解QAT的四个核心：模拟量化、伪量化节点、STE，以及完整的训练流程。

1. 模拟量化：让模型提前“体验”低精度生活

模拟量化，英文叫 Fake Quantization。名字很直白——它不是真的把模型转成int8，而是在浮点计算中模拟量化效果。

具体怎么做？前向传播时，先把浮点值量化成int8，再反量化回浮点。这样模型看到的虽然是浮点数，但已经带上了量化误差。我刚开始接触这个概念时，觉得有点绕：既然最终要量化，为什么不直接训int8？

原因很简单：梯度下降需要连续可导。int8的取整操作不可导，没法直接训练。所以模拟量化就是“中间人”——前向走量化流程，反向走浮点梯度更新。

模拟量化的数学表达：

# 伪代码：模拟量化过程
def fake_quantize(x, scale, zero_point, num_bits=8):
    # 量化
    x_int = torch.round(x / scale + zero_point)
    # 截断到 [0, 2^num_bits - 1]
    x_int = torch.clamp(x_int, 0, 2**num_bits - 1)
    # 反量化
    x_fake = (x_int - zero_point) * scale
    return x_fake

你看，输入x是浮点，输出x_fake也是浮点，但中间经历了量化截断。这个截断误差，就是模型需要去“适应”的东西。

2. 伪量化节点：量化过程的“监控摄像头”

伪量化节点（FakeQuantize Node）是模拟量化的具体实现。在PyTorch中，它就是torch.quantization.FakeQuantize这个模块。

这些节点会被插入到模型的权重和激活值后面。我习惯在卷积层、全连接层后面都加上。它们的作用有两个：

记录量化参数：scale和zero_point会在训练中通过移动平均更新
执行模拟量化：前向时做量化+反量化，反向时传递梯度

嗯，这里要注意：伪量化节点只在训练时生效。推理时，它们会被移除，替换成真正的量化算子。

我的经验： 伪量化节点插入的位置很关键。只量化权重不量化激活，精度可能不够；全量化又可能训练不稳定。我一般先量化权重，再逐步加入激活量化，观察loss曲线。

3. 直通估计器（STE）：解决“不可导”问题的神来之笔

这是QAT最核心的trick。量化中的round操作，梯度几乎处处为0。你想想看，梯度为0，网络还怎么学？

STE（Straight-Through Estimator）的做法很粗暴：前向传播用真实的量化函数，反向传播假装量化函数是恒等映射。也就是说，梯度直接“穿透”量化节点，原封不动地传回去。

数学上就是：

# 前向：y = quantize(x)
# 反向：∂Loss/∂x ≈ ∂Loss/∂y  （忽略量化带来的梯度变化）

我曾经在一个语音模型上试过，不加STE，loss直接炸了；加上STE，训练稳稳的。当然，STE是个近似，它假设量化误差对梯度的影响可以忽略。实践中这个假设基本成立，因为量化误差相对于梯度来说是小量。

避坑指南： 我曾经在某个项目中，把STE的梯度直通写成了y = x + (quantize(x) - x).detach()。这个写法虽然等价，但容易让人误解。记住：detach()是切断梯度，而STE需要保留梯度。正确做法是直接用torch.round，然后依赖STE的默认行为。

4. QAT训练流程：从浮点到量化就绪

完整的QAT流程，我总结为三步走：

预训练浮点模型：先用float32训一个高精度baseline。这一步不能省，QAT不是从零训练，而是微调。
插入伪量化节点：在模型所有需要量化的地方（权重、激活）插入FakeQuantize。我建议先用默认的对称量化，跑通后再调非对称。
微调训练：用较小的学习率（通常是预训练的1/10）继续训练几个epoch。观察loss收敛后，导出量化模型。

这里有个细节：微调时，BN层需要特殊处理。因为伪量化节点改变了激活值的分布，BN的running mean/std需要重新计算。我一般会在微调初期冻结BN，等量化参数稳定后再解冻。

PyTorch QAT 示例代码片段：

import torch
import torch.quantization as quant

# 1. 加载预训练模型
model = torch.load('float_model.pth')

# 2. 设置为QAT模式
model.qconfig = quant.get_default_qat_qconfig('fbgemm')
quant.prepare_qat(model, inplace=True)

# 3. 微调
optimizer = torch.optim.SGD(model.parameters(), lr=1e-4)
for epoch in range(5):
    for data, target in dataloader:
        optimizer.zero_grad()
        output = model(data)
        loss = criterion(output, target)
        loss.backward()
        optimizer.step()

# 4. 转换为量化模型
model.eval()
model = quant.convert(model, inplace=True)

你看，代码其实很简洁。但背后的原理——模拟量化、STE、伪量化节点——才是真正需要理解的东西。我见过不少同学直接调API，精度掉了也不知道怎么调，就是因为没搞懂这些底层机制。

最后说一句：QAT不是万能的。如果模型本身过拟合严重，QAT反而可能放大量化误差。这时候先做正则化，再上QAT，效果会好很多。

📚 课程目录

量化感知训练（QAT）原理：模拟量化、伪量化节点、直通估计器（STE）、QAT训练流程

1. 模拟量化：让模型提前“体验”低精度生活

2. 伪量化节点：量化过程的“监控摄像头”

3. 直通估计器（STE）：解决“不可导”问题的神来之笔

4. QAT训练流程：从浮点到量化就绪