3. 量化感知训练(QAT)：原理介绍，如何在YOLOv5中插入伪量化节点

好，咱们接着聊量化。上一节讲了PTQ，也就是训练后量化。说白了就是模型训练完了，咱们再想办法给它“瘦身”。但PTQ有个问题——精度损失有时候挺明显的，尤其是对YOLO这种检测模型，边界框稍微偏一点，结果就差很多。

那怎么办？这时候就该QAT上场了。量化感知训练，名字听着挺唬人，其实核心思想很简单：让模型在训练阶段就“适应”量化带来的误差。就像你练投篮，如果一开始就戴着负重手套练，等比赛时摘掉手套，手感反而更好。QAT就是这个道理。

3.1 QAT的核心原理

先说说QAT到底在干什么。我个人习惯把它拆成三步理解：

  1. 插入伪量化节点——在模型的关键位置（比如卷积层前后）插入一些“模拟器”
  2. 前向传播模拟量化——训练时，这些节点会模拟量化的舍入误差
  3. 反向传播正常更新——梯度计算时绕过这些节点，权重照样更新

这里有个关键点：伪量化节点只在前向传播时生效。反向传播时，我们用的是“直通估计器”（STE），说白了就是假装这些节点不存在，让梯度直接穿过去。为什么这么做？因为量化操作（比如四舍五入）的梯度是0，没法训练。STE就是绕开这个坑。

核心公式（简化版）：

前向：Q(x) = round(clamp(x / scale) + zero_point)

反向：∂Q/∂x ≈ 1 （STE近似）

嗯，这里要注意：STE只是个近似，不是精确的。我在项目中遇到过，如果量化位宽太低（比如4bit），STE的近似误差会累积，导致模型不收敛。所以一般建议从8bit开始试。

3.2 为什么YOLOv5需要QAT？

你想想看，YOLOv5的检测头里有很多小卷积层，这些层对数值精度特别敏感。PTQ做完后，我经常发现mAP掉了2-3个点，尤其是小目标检测，几乎废了一半。

QAT的好处在于：

  • 精度恢复——通常能把掉点控制在0.5%以内
  • 部署友好——训练完直接导出量化模型，不用再校准
  • 硬件适配——很多NPU、TPU都支持QAT导出的模型

我曾经在一个边缘设备项目里，用PTQ量化后模型精度掉了4.2%，客户直接拒收。后来换成QAT，只掉了0.3%，一次过。所以你说QAT重不重要？

3.3 在YOLOv5中插入伪量化节点

好，到了实操环节。YOLOv5本身没有内置QAT支持，但我们可以用PyTorch的torch.quantization模块手动插入。我一般按以下步骤来：

3.3.1 准备工作

首先，确保你的PyTorch版本>=1.8.0。然后导入必要的模块：

import torch
import torch.nn as nn
from torch.quantization import QuantStub, DeQuantStub, default_qconfig
from torch.quantization import prepare_qat, convert

3.3.2 修改YOLOv5模型结构

我们需要在模型的前向传播中插入QuantStub和DeQuantStub。说白了，就是在输入和输出处加两个“转换器”。

class QuantizedYOLOv5(nn.Module):
    def __init__(self, base_model):
        super().__init__()
        self.quant = QuantStub()  # 输入量化
        self.model = base_model
        self.dequant = DeQuantStub()  # 输出反量化

    def forward(self, x):
        x = self.quant(x)  # 模拟量化输入
        x = self.model(x)
        x = self.dequant(x)  # 恢复浮点输出
        return x

我的经验：别只加在输入输出上。对于YOLOv5的检测头，我建议在每个卷积层后面也加伪量化节点。这样训练出来的模型，量化误差分布更均匀。

3.3.3 配置量化参数

接下来，设置量化配置。我个人习惯用default_qat_qconfig：

model.qconfig = default_qat_qconfig  # 使用默认QAT配置
# 或者自定义：
# model.qconfig = torch.quantization.get_default_qat_qconfig('fbgemm')

然后调用prepare_qat：

model = prepare_qat(model, inplace=True)

这一步会在模型里插入FakeQuantize节点。你可以用print(model)看看，每个卷积层前后都多了些东西。

3.3.4 微调训练

这一步很关键。QAT不是从头训练，而是在预训练模型基础上微调。我建议：

  • 学习率调低——一般是原学习率的1/10，比如从0.01降到0.001
  • 训练轮数减少——5-10个epoch就够了，太多反而过拟合
  • 冻结BN层——量化对BN层影响大，我一般先冻结再微调
# 冻结BN层
for module in model.modules():
    if isinstance(module, nn.BatchNorm2d):
        module.eval()

# 正常训练循环
for epoch in range(5):
    for images, targets in dataloader:
        outputs = model(images)
        loss = criterion(outputs, targets)
        loss.backward()
        optimizer.step()

注意：训练时一定要用model.train()模式。我见过有人用eval()模式跑QAT，结果伪量化节点不更新，精度一点没提升。

3.3.5 导出量化模型

微调完成后，用convert把伪量化节点变成真正的量化操作：

model = convert(model, inplace=True)
# 现在model已经是量化模型了，可以直接导出为torchscript
traced_model = torch.jit.trace(model, example_input)
traced_model.save('yolov5s_quantized.pt')

嗯，这里要注意：convert之后模型就不能再训练了。所以一定要在确认精度满意后再执行这一步。

3.4 避坑指南

做QAT时，我踩过不少坑。分享几个典型的：

  • 伪量化节点位置不对——我曾经把节点插在激活函数后面，结果精度反而更差。正确做法是插在卷积层前后，而不是激活函数前后。
  • 学习率没调——直接用原学习率跑QAT，loss直接飞了。后来改成1/10，稳得很。
  • 数据集太小——QAT需要少量数据做校准，但数据太少（比如不到100张）会导致量化参数不准。我一般至少用500张。

总结一下QAT的流程：

  1. 加载预训练YOLOv5模型
  2. 插入QuantStub/DeQuantStub
  3. 配置qconfig并调用prepare_qat
  4. 用低学习率微调5-10个epoch
  5. 调用convert导出量化模型

说实话，QAT比PTQ麻烦一些，但效果确实好。尤其是你要把YOLOv5部署到嵌入式设备上，QAT几乎是必经之路。下一节我会讲怎么把量化后的模型部署到TensorRT和OpenVINO上，到时候你就知道QAT的威力了。

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