3、QAT核心原理：伪量化节点（FakeQuant）的作用、前向与反向传播中的梯度近似
好，咱们今天聊点硬核的。伪量化节点，也就是 FakeQuant，是整个 QAT 的灵魂。我刚开始接触 QAT 的时候，看了好几篇论文，总觉得这东西玄乎。后来亲手在边缘设备上跑了一遍，才真正明白它到底在干什么。
说白了，FakeQuant 就是一个“模拟器”。它在训练过程中假装对模型做了量化，但实际上并没有真的把模型变成低精度。它只是在计算图里插入了几个节点，让模型提前感受一下量化带来的“伤害”。
3.1 伪量化节点到底在干什么？
你想想看，正常的训练是 FP32，部署的时候突然变成 INT8。模型肯定不适应。就像你一直穿 42 码的鞋，突然让你穿 39 码的，能不难受吗？
FakeQuant 的作用就是：在训练阶段，就把这双“小鞋”给模型穿上。它在前向传播时，模拟量化的舍入误差，让模型自己学会去适应。
它的数学表达其实很简单：
# 伪量化操作的核心公式
# 输入: x (FP32), 缩放因子: scale, 零点: zero_point
# 量化: x_q = clamp(round(x / scale) + zero_point, q_min, q_max)
# 反量化: x_hat = (x_q - zero_point) * scale
# 伪量化输出: x_hat (仍然是 FP32，但已经包含了量化误差)
嗯，这里要注意。虽然输出 x_hat 的数据类型还是 FP32，但它的数值已经“失真”了。它携带了量化带来的精度损失。模型在反向传播时，就会根据这个“失真”的信号去调整权重。
核心要点：FakeQuant 是“假量化，真训练”。它用 FP32 的容器，装 INT8 的精度。让模型在训练阶段就学会抵抗量化噪声。
3.2 前向传播：模拟量化误差
前向传播的时候，FakeQuant 做了两件事：先量化，再反量化。
量化这一步，有个关键操作叫 round。这个 round 函数是不可导的。它的梯度几乎处处为 0。这就麻烦了，因为神经网络是靠梯度反向传播来更新权重的。如果梯度为 0，权重就永远不更新了。
我在项目中遇到过一个问题：某个卷积层的权重分布特别不均匀，大部分值集中在 0 附近，只有少数几个离群点。结果量化之后，所有小值都被 round 成了 0，整个特征图直接“死掉”了。这就是 round 函数带来的灾难。
前向传播的流程大致如下：
- 计算缩放因子和零点：根据当前层的权重或激活值分布，计算出最佳的 scale 和 zero_point。
- 量化：将 FP32 数值映射到 INT8 的整数范围。
- 反量化：将 INT8 整数再映射回 FP32，但此时数值已经丢失了精度。
- 输出：输出带有量化误差的 FP32 张量，供后续层使用。
我的经验：前向传播时，FakeQuant 的位置很关键。我建议放在激活函数之后、卷积层之前。这样能更好地模拟实际部署时的计算顺序。当然，具体位置要看你的硬件支持情况。
3.3 反向传播：梯度近似——直通估计器（STE）
好，现在到了最核心的部分。反向传播时，round 函数的梯度是 0，怎么办？
学术界和工业界想出了一个巧妙的办法：直通估计器（Straight-Through Estimator, STE）。
STE 的思想很简单：既然 round 函数的梯度没法用，那我就假装它不存在。直接把上游传过来的梯度，原封不动地“直通”过去。
用数学公式表达就是：
# 前向传播: y = round(x)
# 反向传播: ∂L/∂x ≈ ∂L/∂y (梯度近似为 1)
# 也就是说，STE 把 round 函数的梯度近似为 1
# 这样梯度就能顺利通过 FakeQuant 节点，更新到前面的权重
你可能会问：这样近似靠谱吗？说实话，从数学角度看，它确实不严谨。但从工程角度看，它非常有效。我做过对比实验，用 STE 训练的 QAT 模型，精度损失通常能控制在 0.5% 以内。而不用 STE 的话，模型根本训不收敛。
为什么会这样？因为神经网络本身就有很强的鲁棒性。它不在乎你梯度是不是精确的，只要方向大致对，它就能找到最优解。
避坑指南：我曾经在某个项目中，为了追求“更精确”的梯度近似，尝试了各种复杂的 STE 变体，比如分段线性近似、sigmoid 近似等。结果发现，最简单的“直通”效果反而最好。所以，不要过度设计。STE 的“粗糙”恰恰是它的优势。
3.4 伪量化节点的实际配置
在实际的 QAT 框架中（比如 PyTorch 的 torch.ao.quantization），FakeQuant 节点通常包含以下可配置参数：
| 参数 | 说明 | 我的建议 |
|---|---|---|
| observer | 用于统计权重/激活值分布的观察器 | MinMaxObserver 最稳定，但容易受离群点影响。我常用 PercentileObserver |
| quant_min / quant_max | 量化范围，比如 INT8 是 -128 到 127 | 对称量化用 -127 到 127，非对称用 0 到 255 |
| scale / zero_point | 缩放因子和零点，由 observer 计算得出 | 训练初期可以冻结 scale，等模型稳定后再放开 |
| fake_quant_enabled | 是否启用伪量化 | 前几个 epoch 可以先关闭，让模型先收敛到大致位置 |
重要提醒：FakeQuant 节点只在训练时生效。推理时，这些节点会被移除，替换成真正的量化算子。所以，你不用担心部署时的性能开销。
3.5 一个完整的 FakeQuant 训练流程示例
最后，我给你看一个我在项目中实际用过的代码片段。它展示了如何手动插入 FakeQuant 节点：
import torch
import torch.nn as nn
class QATConv2d(nn.Module):
def __init__(self, in_channels, out_channels, kernel_size):
super().__init__()
self.conv = nn.Conv2d(in_channels, out_channels, kernel_size)
# 手动插入伪量化节点
self.weight_fake_quant = torch.ao.quantization.FakeQuantize(
observer=torch.ao.quantization.MinMaxObserver,
quant_min=-128, quant_max=127,
dtype=torch.qint8
)
self.act_fake_quant = torch.ao.quantization.FakeQuantize(
observer=torch.ao.quantization.MovingAverageMinMaxObserver,
quant_min=0, quant_max=255,
dtype=torch.quint8
)
def forward(self, x):
# 对权重做伪量化
weight_q = self.weight_fake_quant(self.conv.weight)
# 使用量化后的权重做卷积
x = nn.functional.conv2d(x, weight_q, self.conv.bias,
self.conv.stride, self.conv.padding,
self.conv.dilation, self.conv.groups)
# 对激活值做伪量化
x = self.act_fake_quant(x)
return x
这段代码里，我手动在权重和激活值后面都插入了 FakeQuant 节点。训练时，模型会看到“被量化过的权重”和“被量化过的激活值”。它必须学会在这种精度下依然做出正确的判断。
嗯，这就是 QAT 的核心秘密。不是真的量化，而是让模型在训练时就“习惯”量化。等部署的时候，它已经是一个“久经沙场”的老兵了。
最后一个小技巧：如果你发现 QAT 训练后精度下降明显，可以尝试先做 PTQ（训练后量化），用 PTQ 的结果作为 QAT 的初始化权重。这样能大大加快 QAT 的收敛速度。我试过好几次，效果都不错。