3. 量化感知训练（QAT）

大家好，我是你们的老朋友。今天我们来聊聊量化里最硬核的一块——量化感知训练（QAT）。

说实话，很多人一开始接触量化，都是从PTQ（训练后量化）入门的。毕竟简单，模型训完直接转，省事。但PTQ有个问题——精度掉得厉害，尤其是小模型或者对精度敏感的任务。这时候，QAT就派上用场了。

QAT的核心思想很简单：让模型在训练阶段就“体验”量化带来的误差，然后通过微调去适应它。说白了，就是让模型学会在低精度下“生存”。

3.1 QAT的原理：模拟量化，反向传播

QAT的原理，我可以用一句话概括：在前向传播时模拟量化，在反向传播时假装没量化。

为什么会这样？因为量化本身是个离散操作，比如把浮点数四舍五入到整数，这个操作不可导。如果直接对量化后的模型做反向传播，梯度传不回去。所以QAT用了一个技巧——直通估计器（STE，Straight-Through Estimator）。

STE的做法是：前向传播时，我们真的把权重和激活值量化到低精度（比如INT8），计算损失；反向传播时，我们假装量化操作不存在，直接把梯度传回给浮点权重。这样，浮点权重就能根据量化后的误差更新了。

嗯，这里要注意：更新的是浮点权重，不是量化后的整数权重。量化后的整数权重只是前向传播时“临时”用的。

我在项目中遇到过一个问题：如果学习率设置太大，浮点权重更新幅度过大，会导致量化后的整数权重频繁跳变，模型反而训不收敛。所以QAT的学习率通常要比普通训练小一个数量级。

3.2 伪量化节点（FakeQuant）：QAT的核心组件

QAT的实现，靠的是在计算图中插入伪量化节点（FakeQuant）。这个节点干了两件事：

• 量化：把浮点数映射到低精度整数（比如INT8）
• 反量化：再把整数映射回浮点数（但精度已经损失了）

所以，伪量化节点的输出仍然是浮点数，但数值已经“模拟”了量化后的精度损失。模型看到的是浮点数，但感受到的是量化误差。

伪量化节点的数学表达是这样的：

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

你看，输入是浮点，输出也是浮点，但中间经过了量化-反量化的过程。这就是“伪量化”名字的由来——看起来没变，实际上精度已经掉了。

我个人习惯在模型的权重和激活值后面都插入FakeQuant节点。权重量化通常用per-tensor（整个张量共享一个scale），激活值量化用per-tensor或者per-channel，看具体硬件支持。

3.3 QAT的完整流程：从浮点到INT8

QAT的流程，我总结为四步走：

1. 准备预训练模型：先拿一个浮点模型，在原始任务上训好。这一步不能省，QAT不是从零开始训练。
2. 插入FakeQuant节点：在模型的权重和激活值后面插入伪量化节点。注意，BN层（Batch Normalization）的统计量需要冻结，否则QAT过程中BN的均值和方差会乱掉。
3. 微调训练：用较小的学习率（比如1e-5）继续训练几个epoch。数据量不需要太大，一般一个epoch就够了，但要看任务复杂度。
4. 导出量化模型：训练完成后，把FakeQuant节点移除，导出真正的INT8模型（权重是INT8，激活值也是INT8）。

我曾经踩过一个坑：在第二步插入FakeQuant节点时，没有冻结BN层。结果训练过程中BN的running_mean和running_var不断更新，导致量化后的模型精度比PTQ还差。后来我学乖了，插入FakeQuant后第一件事就是冻结BN。

下面是一个简化的QAT训练代码示例：

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

# 1. 准备预训练模型
model = MyModel()
model.load_state_dict(torch.load('pretrained.pth'))

# 2. 插入FakeQuant节点
model.qconfig = quant.get_default_qat_qconfig('fbgemm')
quant.prepare_qat(model, inplace=True)

# 冻结BN层
for module in model.modules():
    if isinstance(module, nn.BatchNorm2d):
        module.eval()  # 冻结BN统计量

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

# 4. 导出量化模型
model.eval()
quantized_model = quant.convert(model, inplace=False)
torch.save(quantized_model.state_dict(), 'qat_model.pth')

注意，quant.convert这一步会把FakeQuant节点移除，替换成真正的量化算子。转换后的模型，权重就是INT8了。

3.4 QAT vs PTQ：什么时候该用谁？

这个问题，我经常被问到。我的建议是：能PTQ就别QAT，除非PTQ精度不够。

为什么？因为QAT需要训练，需要数据，需要调参，成本高。PTQ只需要一个校准集，跑一遍前向传播就完事了。

但PTQ有个硬伤：它对小模型和敏感任务不友好。比如MobileNet、BERT这类模型，PTQ后精度可能掉2-3个点，甚至更多。这时候QAT就香了。

我整理了一个对比表格，方便你决策：

对比维度	PTQ（训练后量化）	QAT（量化感知训练）
是否需要训练	不需要	需要微调
是否需要数据	少量校准集（几百张图）	训练集（几千张图以上）
精度损失	大模型0.5-1%，小模型2-5%	通常0.1-0.5%
部署成本	低，一步到位	高，需要训练+转换
适用场景	大模型、精度不敏感任务	小模型、精度敏感任务

我个人经验是：如果模型参数量超过100M，PTQ通常够用；如果模型小于10M，建议直接上QAT。中间那档，先试PTQ，精度不行再换QAT。

3.5 QAT的进阶技巧：如何训得更好？

QAT虽然效果好，但也不是无脑训就能成的。我分享几个实战技巧：

• 学习率衰减：QAT的学习率建议从1e-5开始，每2个epoch衰减0.5倍。训太快容易震荡。
• 逐层量化：如果模型很大，可以逐层插入FakeQuant节点，先量化前几层，再量化后几层。这样梯度更稳定。
• 混合精度：有些层对量化特别敏感（比如第一层卷积），可以保留FP16或FP32，其他层用INT8。这叫混合精度量化。
• 校准集选择：QAT用的数据最好和训练集分布一致。如果用了分布不同的数据，模型会“学歪”。

我记得有一次，一个语音识别模型QAT后精度反而下降了。排查了半天，发现是校准集里混了噪声数据。换了干净数据后，精度立刻回升。所以，数据质量比数量更重要。

3.6 本章小结：QAT的核心要点

好了，我们来捋一捋QAT的核心要点：

• 原理：用STE解决量化不可导问题，让模型在训练时适应量化误差
• 组件：FakeQuant节点，模拟量化-反量化过程
• 流程：预训练 → 插入FakeQuant → 微调 → 导出INT8模型
• 对比：QAT精度高但成本高，PTQ成本低但精度可能不够

QAT不是银弹，但它是目前精度最高的量化方案。如果你的模型对精度要求苛刻，或者PTQ后精度掉得没法看，那就试试QAT吧。

最后，送大家一句话：量化不是目的，部署才是。选对方案，少走弯路。