3、量化感知训练：让模型学会适应低精度

说实话，很多人刚开始接触模型量化时，都会有一个疑问：

「为什么直接做后训练量化，精度掉得那么厉害？」

我当年第一次跑 PTQ 的时候，一个分类模型从 FP32 的 92% 直接掉到 85%。当时我就懵了。后来才发现，模型根本不知道你要把它变成低精度。它所有的权重和激活值，都是为高精度计算精心调出来的。你突然一刀切下去，它当然会「水土不服」。

那怎么办？

很简单——让模型在训练阶段就「知道」自己将来要跑在低精度下。这就是量化感知训练，简称 QAT。

3.1 核心思想：模拟量化误差

QAT 的原理，说白了就是四个字：假装量化。

在训练的前向传播过程中，我们插入一些假的量化节点。这些节点会把浮点数先量化成整数，再反量化回浮点数。这样模型在前向计算时，就已经感受到了量化带来的精度损失。

反向传播的时候呢？梯度还是用浮点数来算。因为量化操作本身不可导，我们得用直通估计器（STE）来绕过这个坎。

关键点：

  • 前向：模拟量化（量化 + 反量化）
  • 反向：梯度直通（忽略量化节点的梯度）
  • 权重更新：用浮点数梯度更新浮点数权重

我习惯把 QAT 比作「戴着镣铐跳舞」。模型在训练时就被限制了精度，它必须学会在这种约束下依然做出正确的判断。等训练结束，我们把那些模拟节点一摘，换上真正的量化算子，精度损失就小得多了。

3.2 伪量化节点：FakeQuant

QAT 里最核心的操作，就是 FakeQuant。它长什么样？

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

你看，输入是浮点数，输出也是浮点数。但中间走了一圈整数，精度已经丢了。模型在训练时，每次前向都要经历这么一遭。

我在项目中遇到过一个问题：scale 和 zero_point 怎么更新？

有两种做法：

  • 基于统计：每次前向时，根据当前 batch 的数据分布重新计算 scale 和 zero_point。这种方法更灵活，但训练不稳定。
  • 基于学习：把 scale 和 zero_point 当成可训练参数，用梯度去更新。我个人的经验是，这种方法收敛更慢，但最终精度往往更高。

我的建议：

刚开始做 QAT 时，先用基于统计的方法。等模型跑通了，再尝试学习型的 scale。别一上来就搞复杂的，容易翻车。

3.3 训练策略：从 PTQ 到 QAT 的过渡

QAT 不是从头开始训练的。你想想看，一个随机初始化的模型，你让它一边学任务一边学量化，太难了。

正确的做法是：

  1. 先训练一个 FP32 的基线模型，精度尽可能高。
  2. 做一次 PTQ，看看精度掉多少。如果掉得不多（比如 1% 以内），其实没必要上 QAT。
  3. 加载 FP32 权重，插入 FakeQuant 节点，用小学习率微调。
  4. 微调结束后，导出为 INT8 模型。

我曾经踩过一个坑：微调时学习率设得太大，结果模型直接跑飞了。后来我学乖了，一般用 FP32 训练时学习率的 1/10 到 1/100。而且只微调 10% 到 20% 的 epoch 就够了。

注意：

QAT 微调时，BN 层的统计量要小心处理。我建议冻结 BN 层的 running mean 和 running var，只更新权重。否则训练不稳定，精度反而下降。

3.4 代码实战：PyTorch 中的 QAT

PyTorch 对 QAT 的支持已经很成熟了。我们来看一个实际例子。

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

# 1. 定义模型
class MyModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = nn.Conv2d(3, 64, 3)
        self.relu = nn.ReLU()
        self.fc = nn.Linear(64, 10)
    
    def forward(self, x):
        x = self.conv1(x)
        x = self.relu(x)
        x = x.mean([2, 3])
        x = self.fc(x)
        return x

model = MyModel()
model.load_state_dict(torch.load('fp32_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 train_loader:
        optimizer.zero_grad()
        output = model(data)
        loss = nn.CrossEntropyLoss()(output, target)
        loss.backward()
        optimizer.step()

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

这段代码里，prepare_qat 会自动在卷积层和全连接层前后插入 FakeQuant 节点。微调完成后，convert 会把 FakeQuant 节点替换成真正的量化算子。

嗯，这里要注意：convert 之前一定要调用 model.eval()。我刚开始做的时候忘了这步，结果导出的模型精度一塌糊涂。后来查了半天才发现是 BN 层的问题。

3.5 精度与速度的权衡

QAT 不是万能的。它也有自己的代价。

方法 精度损失 训练时间 部署难度
PTQ（后训练量化） 1% - 5% 无 低
QAT（量化感知训练） 0.1% - 1% 增加 10% - 30% 中
FP32 基线 0% 正常 高（推理慢）

从表格能看出来，QAT 的精度损失确实小，但代价是训练时间变长。我个人习惯是：

  • 如果 PTQ 精度损失在 1% 以内，直接用 PTQ，省事。
  • 如果损失在 1% - 3%，上 QAT 微调。
  • 如果损失超过 3%，先检查模型结构，是不是某些层对量化特别敏感。

避坑指南：

我曾经在一个语义分割模型上做 QAT，精度死活上不去。后来发现是 skip connection 上的加法操作对量化太敏感。解决办法是在加法之前单独加一个 FakeQuant 节点，把两路输入的尺度对齐。精度一下子就回来了。

3.6 什么时候该用 QAT？

说实话，不是所有场景都需要 QAT。我总结了几条经验：

  • 小模型优先用 QAT：MobileNet、SqueezeNet 这类轻量模型，对量化更敏感，PTQ 往往掉精度严重。
  • 大模型可以先用 PTQ 试试：ResNet50 以上的大模型，冗余度高，PTQ 效果往往不错。
  • 检测、分割模型建议 QAT：分类任务对量化容忍度高，但检测和分割的回归头很脆弱，QAT 几乎是必须的。
  • 如果部署平台不支持混合精度：有些芯片只支持对称量化或特定位宽，QAT 可以提前适配。

你想想看，QAT 的本质是什么？是让模型学会在低精度下生存。它不是在「修复」量化带来的问题，而是在「预防」这些问题。这种思路，其实在很多工程领域都适用——提前暴露问题，总比事后补救要好。

好了，这一章就到这里。下一章我们会聊聊量化推理引擎的选型，看看不同硬件上怎么跑量化模型才最快。

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