量化感知训练（QAT）入门

各位同学，今天我们来聊聊量化感知训练。说实话，很多刚入行的朋友会把QAT想得很神秘，觉得是什么黑科技。其实说白了，QAT就是让模型在训练过程中就学会适应低精度推理的一种方法。

我个人习惯把QAT比作「带着镣铐跳舞」。你想想看，模型本来在FP32的世界里自由自在，突然要把它塞进INT8的小盒子里，它肯定不适应。QAT就是提前告诉模型：嘿，你以后要在小盒子里生活，现在先练练。

QAT的核心原理

QAT的原理其实不复杂。我们在训练过程中插入一些特殊的节点——伪量化节点（Fake Quantize Nodes）。这些节点会模拟量化-反量化的过程，让模型的前向传播感受到量化带来的精度损失。

为什么叫「伪量化」？因为它只模拟，不真正做量化。具体来说：

• 前向传播：输入先被量化到INT8，再反量化回FP32，然后继续计算
• 反向传播：梯度直接跳过伪量化节点，用直通估计器（STE）来近似梯度

我在项目中遇到过一个问题：刚开始用QAT时，总觉得训练完的模型精度还不如直接PTQ。后来才发现，原来是学习率没调好。嗯，这里要注意，QAT的学习率一般要比正常训练小一个数量级。

Q(x) = round(clamp(x / scale + zero_point, qmin, qmax))
DQ(x) = (Q(x) - zero_point) * scale

在PyTorch中插入伪量化节点

PyTorch提供了非常方便的QAT工具。我个人推荐使用torch.quantization.QuantStub和torch.quantization.DeQuantStub来标记量化边界。

来看一个简单的例子：

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

class QATModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.quant = quant.QuantStub()
        self.conv1 = nn.Conv2d(3, 64, 3)
        self.relu = nn.ReLU()
        self.dequant = quant.DeQuantStub()
    
    def forward(self, x):
        x = self.quant(x)      # 插入伪量化节点
        x = self.conv1(x)
        x = self.relu(x)
        x = self.dequant(x)    # 反量化
        return x

# 准备QAT
model = QATModel()
model.qconfig = quant.get_default_qat_qconfig('fbgemm')
quant.prepare_qat(model, inplace=True)

# 正常训练...
# 训练完成后转换
model.eval()
quant.convert(model, inplace=True)

这里有个坑，我曾经踩过：prepare_qat和convert的顺序不能搞反。先prepare_qat训练，再convert做实际量化。如果你先convert再训练，那伪量化节点就没了，QAT等于白做。

训练技巧与注意事项

QAT的训练技巧，我总结了几条实战经验：

QAT的适用场景

不是所有模型都适合QAT。我个人觉得，以下场景用QAT效果最好：

• 小模型：MobileNet、SqueezeNet这类轻量网络，PTQ掉点严重，QAT能挽回不少
• 低比特量化：INT4甚至INT2量化，必须用QAT
• 精度要求高：比如人脸识别、医疗影像，掉点0.5%都不能接受

反过来，如果你的模型很大（比如ResNet-152），或者量化位宽够宽（INT8），PTQ可能就够用了。没必要为了QAT而QAT，毕竟训练时间也是成本。

知识体系结构图

下面这张图展示了QAT的完整知识体系，我把它画成了流程图，方便你理解各个环节的关系：

好了，关于QAT入门就讲到这里。记住一句话：QAT不是万能的，但没有QAT是万万不能的。尤其是在嵌入式设备上，每一分精度都弥足珍贵。