3、量化感知训练：如何让模型适应低精度？—— 模拟量化、伪量化节点、BN层融合技巧

好，咱们接着聊量化。上一节讲了训练后量化（PTQ），说白了就是模型训练完了，我们直接把它转成低精度。这个方法快，但精度损失有时候让人头疼。

那有没有办法让模型从一开始就“习惯”低精度呢？

有。这就是量化感知训练（QAT）。

我个人习惯把QAT比作“带着镣铐跳舞”。你让模型在训练时就模拟低精度的环境，它自然就学会了在这种约束下表现好。今天我们就来拆解QAT的三个核心：模拟量化、伪量化节点，还有那个让人又爱又恨的BN层融合。

3.1 模拟量化：让模型提前“体验”低精度

先问一个问题：为什么训练后量化会掉精度？

说白了，模型在训练时用的是FP32，权重和激活值都是连续的、高精度的。你突然把它变成INT8，相当于让一个习惯了开跑车的人去开拖拉机——他肯定不适应。

模拟量化的思路很简单：在训练过程中，就假装我们已经把模型量化了。

具体怎么做呢？在前向传播时，我们把权重和激活值“模拟”成低精度的样子。比如，先量化成INT8，再反量化回FP32。这样模型看到的虽然是FP32的数值，但已经带上了量化误差。

我在项目中遇到过这种情况：一个分类模型，PTQ之后准确率从92%掉到了88%。用了QAT之后，直接拉回到91.5%。效果很明显。

模拟量化的数学表达其实不复杂：

# 模拟量化的伪代码
def simulate_quantize(x, scale, zero_point, num_bits=8):
    # 1. 量化
    x_int = torch.round(x / scale + zero_point)
    # 2. 截断到有效范围
    x_int = torch.clamp(x_int, 0, 2**num_bits - 1)
    # 3. 反量化回浮点
    x_fp = (x_int - zero_point) * scale
    return x_fp

你看，前向传播时，数据走了一遍“量化-反量化”的流程。反向传播时呢？这里有个关键点——round操作是不可导的。怎么办？

嗯，这里要用到“直通估计器”（STE）。简单说，就是前向传播用round，反向传播时假装round不存在，直接把梯度传过去。这招虽然粗糙，但实践中效果很好。

核心要点：模拟量化不是真的量化，而是让模型在训练时“感受”量化误差，从而学会抵抗这种误差。

3.2 伪量化节点：QAT的“心脏”

刚才说的模拟量化，在代码里是怎么实现的？

答案就是——伪量化节点（FakeQuantize Node）。

你想想看，如果要在整个计算图中插入量化操作，总不能手动改每一层吧？伪量化节点就是干这个的。它是一个特殊的算子，专门负责“模拟量化”。

在PyTorch中，torch.quantization.FakeQuantize就是干这个的。它通常被插入在卷积层、全连接层之后，或者激活函数之前。

我记得有一次调试一个模型，发现量化后精度崩了。排查了半天，原来是某个伪量化节点的位置放错了。它被放在了ReLU之前，导致负值被截断的方式不对。

伪量化节点的工作流程是这样的：

  1. 收集统计信息：在训练过程中，它会观察输入数据的范围（min/max）。
  2. 计算量化参数：根据观察到的范围，计算scale和zero_point。
  3. 执行模拟量化：对输入进行量化-反量化操作。
  4. 更新参数：量化参数本身也是可学习的（有些实现中）。

这里有个细节要注意：伪量化节点在训练时是“活跃”的，但在推理时会被移除，替换成真正的量化算子。所以，QAT训练好的模型，导出时是INT8的，推理速度会快很多。

我的建议：在训练初期，可以先关闭伪量化节点（用observer模式），等模型收敛一段时间后再开启。这样模型先学到好的特征，再适应量化误差，效果更稳定。

3.3 BN层融合技巧：别让BN“捣乱”

好，接下来聊一个实战中特别容易踩坑的地方——BN层。

为什么BN层在QAT中是个问题？

你想想看，BN层在训练时是逐样本归一化的，它依赖于当前batch的均值和方差。但量化时，我们希望每个通道的scale是固定的。如果BN层的统计信息一直在变，量化参数就不好确定。

我曾经在一个项目里，QAT训练了100个epoch，精度死活上不去。后来发现，BN层的running_mean和running_var还在剧烈波动。说白了，模型根本没稳定下来。

解决方案是什么？

BN层融合（BatchNorm Fusion）。

具体做法是：在QAT训练的最后阶段，把BN层的参数“吸收”到卷积层或全连接层中。这样，推理时就不需要单独计算BN了，量化也变得更稳定。

融合的数学原理其实很简单：

# BN层融合的数学推导
# 卷积层输出: y = W * x + b
# BN层: z = gamma * (y - mean) / sqrt(var + eps) + beta
# 融合后: z = (gamma * W / sqrt(var + eps)) * x + (gamma * (b - mean) / sqrt(var + eps) + beta)

你看，融合后的权重和偏置可以直接计算出来。这样，模型就变成了“卷积+激活”的简单结构，量化起来方便多了。

在PyTorch中，可以用torch.quantization.fuse_modules来自动完成这个操作：

import torch.quantization as quant

# 定义要融合的模块
modules_to_fuse = [['conv1', 'bn1', 'relu1'],
                   ['conv2', 'bn2', 'relu2']]

# 执行融合
model = quant.fuse_modules(model, modules_to_fuse)
注意：BN层融合一定要在QAT训练的最后阶段做。如果在训练一开始就融合，BN层的统计信息就无法正常更新了，模型精度会受影响。

3.4 QAT的完整流程

好了，我们把上面三个知识点串起来，看看一个完整的QAT流程是什么样的：

  1. 准备预训练模型：先用FP32训练一个高精度的模型。
  2. 插入伪量化节点：在模型中插入FakeQuantize节点，开启模拟量化。
  3. 微调训练：用较小的学习率继续训练几个epoch，让模型适应量化误差。
  4. BN层融合：在训练的最后阶段，把BN层融合到卷积层中。
  5. 导出量化模型：将模型转换为INT8格式，部署到推理引擎中。

这个流程看起来简单，但每个步骤都有细节。我建议你在实践中多尝试不同的学习率和训练时长。一般来说，QAT的微调只需要FP32训练时间的10%-20%，就能达到不错的效果。

3.5 避坑指南

最后，分享几个我踩过的坑：

  • 学习率太大：QAT微调时，学习率一定要小。我曾经用1e-3，结果模型直接发散。后来改成1e-5，效果就好了。
  • 伪量化节点位置不对：有些层（比如Softmax）不适合量化，要手动排除。
  • BN层没融合就导出：导出的模型里还带着BN层，推理时统计信息对不上，精度崩了。
  • 校准数据集太小：虽然QAT不需要校准集，但如果你用PTQ做对比，校准集太小会导致量化参数不准。

嗯，量化感知训练就讲到这里。说白了，QAT就是用一点点训练成本，换取模型在低精度下的高精度。下一节我们会聊更高级的混合精度量化策略，敬请期待。

← 上一章 📖 返回目录 下一章 →