4. 量化感知训练 (QAT)

各位同学，欢迎来到第四章。

前几章我们聊了PTQ（训练后量化），那种方式简单粗暴，适合大部分场景。但如果你追求极致精度，或者模型本身对量化特别敏感——比如某些轻量级网络，PTQ掉点严重——那你就得请出今天的主角：量化感知训练（QAT）。

说白了，QAT就是在训练过程中“假装”模型已经被量化了。让权重和激活值提前适应低比特的“粗糙感”。这样训练出来的模型，量化后精度几乎无损。

我个人习惯把QAT比作“带着镣铐跳舞”。你想想看，模型在训练时就知道自己将来要被量化，所以它学会了在有限的精度下表达自己。这比事后强行压缩要聪明得多。

4.1 模拟量化：让模型提前“感受”低精度

QAT的核心思想很简单：在前向传播时，插入伪量化节点（FakeQuant）。这些节点会模拟量化和反量化的过程，把浮点数“压缩”成INT8再“恢复”回来。

注意，这里只是模拟。反向传播时，梯度仍然用浮点数计算。这就是所谓的“直通估计器（STE）”技巧。

核心公式（伪量化操作）：

# 伪量化：模拟量化-反量化过程
# 输入：浮点数 x，量化参数 scale, zero_point
# 输出：模拟量化后的浮点数

def fake_quantize(x, scale, zero_point, num_bits=8):
    # 1. 量化（浮点 -> 整型）
    x_int = torch.round(x / scale) + zero_point
    # 2. 截断到有效范围
    qmin = 0
    qmax = 2**num_bits - 1
    x_int = torch.clamp(x_int, qmin, qmax)
    # 3. 反量化（整型 -> 浮点）
    x_fp = (x_int - zero_point) * scale
    return x_fp

嗯，这里要注意：torch.round 是不可导的。所以反向传播时，STE会直接把梯度“绕过”这个取整操作，当作恒等映射来处理。虽然理论上有点粗糙，但实践中效果非常好。

我在项目中遇到过一个问题：如果scale设置得太小，量化后的数值会集中在几个bin里，信息丢失严重。后来我加了一个梯度缩放因子，让scale在训练中慢慢调整，效果才稳定下来。

4.2 伪量化节点（FakeQuant）的插入位置

伪量化节点不是随便插的。通常我们只插在权重和激活值上。具体来说：

  • 权重：每个卷积层或全连接层的权重，在参与计算前先经过伪量化。
  • 激活值：每个卷积层或全连接层的输出，在进入下一层前先经过伪量化。

为什么只插这两个地方？因为量化的瓶颈就在权重和激活值。偏置（bias）通常保持32位浮点，不参与量化。

我的经验： 如果你用的是PyTorch，可以直接用 torch.quantization.FakeQuantize 模块。它会自动管理scale和zero_point的更新。我曾经手写过一套伪量化逻辑，结果发现官方实现已经优化得很好了，别重复造轮子。

伪量化节点的插入流程，我画了一张图，你一看就明白：

伪量化节点插入示意图 输入 (FP32) FakeQuant (激活) 卷积 (FP32) FakeQuant (权重) 输出 (FP32) 说明 红色：伪量化节点 绿色：正常计算层 蓝色：数据流 权重和激活 分别插入FakeQuant

4.3 训练后微调策略：别把模型“练废了”

QAT不是从头训练，而是在一个预训练好的FP32模型基础上进行微调。微调的策略很关键，搞不好会把模型“练废了”。

我曾经犯过一个错误：微调时学习率设得太大，结果模型直接发散，精度比PTQ还差。后来我总结了一套经验：

  1. 学习率要小：通常是原始训练学习率的1/10到1/100。我习惯用1e-5起步，观察loss曲线再调整。
  2. 微调轮数要少：一般1-5个epoch就够了。QAT只是让模型适应量化噪声，不是重新学习特征。
  3. 冻结BN层：BatchNorm的统计量在微调时最好固定住。否则量化后的分布会偏移，导致精度下降。

避坑指南： 我曾经在微调时忘了冻结BN层，结果模型在验证集上精度很高，但量化后直接崩了。原因是BN的running_mean和running_var被微调数据带偏了。记住：QAT微调时，BN层用预训练模型的统计量，不要更新。

4.4 代码实战：PyTorch中的QAT实现

光说不练假把式。我们来看一段实际的QAT微调代码。这里以ResNet-18为例：

import torch
import torch.nn as nn
import torch.optim as optim
from torch.quantization import QuantStub, DeQuantStub, FakeQuantize

# 1. 准备模型（预训练FP32）
model = torch.hub.load('pytorch/vision:v0.10.0', 'resnet18', pretrained=True)

# 2. 插入伪量化节点
model.quant = QuantStub()    # 输入量化
model.dequant = DeQuantStub() # 输出反量化

# 3. 设置QAT配置
model.qconfig = torch.quantization.get_default_qat_qconfig('fbgemm')
# 或者自定义：
# model.qconfig = torch.quantization.QConfig(
#     activation=FakeQuantize.with_args(
#         observer=torch.quantization.MinMaxObserver,
#         quant_min=0, quant_max=255,
#         dtype=torch.quint8, qscheme=torch.per_tensor_affine
#     ),
#     weight=FakeQuantize.with_args(
#         observer=torch.quantization.PerChannelMinMaxObserver,
#         quant_min=-128, quant_max=127,
#         dtype=torch.qint8, qscheme=torch.per_channel_symmetric
#     )
# )

# 4. 准备QAT
model.train()
model = torch.quantization.prepare_qat(model, inplace=True)

# 5. 微调（只训练1个epoch）
optimizer = optim.SGD(model.parameters(), lr=1e-5)
criterion = nn.CrossEntropyLoss()

for batch_idx, (data, target) in enumerate(train_loader):
    optimizer.zero_grad()
    output = model(data)
    loss = criterion(output, target)
    loss.backward()
    optimizer.step()
    
    if batch_idx % 100 == 0:
        print(f'Batch {batch_idx}, Loss: {loss.item():.4f}')

# 6. 转换为量化模型
model.eval()
model = torch.quantization.convert(model, inplace=True)

# 现在model已经是INT8量化模型了
print("QAT完成，模型已量化！")

这段代码里，prepare_qat 会自动把FakeQuant节点插入到合适的位置。convert 则会把FakeQuant节点替换成真正的量化操作。

我的建议： 如果你用的是自定义模型，记得手动指定 qconfig。特别是激活值的量化方案，per_tensor_affine 和 per_channel_symmetric 差别很大。我一般对权重用per_channel，对激活用per_tensor，这样精度和性能比较平衡。

4.5 量化参数的学习与校准

QAT过程中，scale和zero_point是怎么确定的？有两种方式：

方式 原理 优缺点
基于观察者（Observer） 在训练过程中，Observer收集数据的统计信息（min/max, 直方图等），然后计算scale和zero_point。 简单稳定，但需要校准数据。适合大部分场景。
基于梯度学习 把scale和zero_point作为可学习参数，通过反向传播更新。 精度可能更高，但训练不稳定，容易发散。我试过几次，效果时好时坏。

我个人习惯先用Observer方式跑一轮，看看量化误差。如果精度不够，再尝试梯度学习方式。不过说实话，大部分场景下Observer方式已经够用了。

4.6 总结与经验分享

QAT的核心就三点：

  • 模拟量化：用FakeQuant节点让模型提前适应低精度。
  • STE梯度：反向传播时绕过取整操作，保持梯度流通。
  • 微调策略：小学习率、少轮数、冻结BN层。

我在多个项目中验证过，QAT相比PTQ，精度平均能提升0.5-1.5个点。特别是对于MobileNet、EfficientNet这类轻量级网络，QAT几乎是必须的。

嗯，最后提醒一句：QAT虽然效果好，但训练时间会变长。如果你的模型对精度不敏感，或者推理部署平台对量化支持很好，PTQ仍然是首选。别为了追求完美而过度工程化。

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