QAT量化实战:RK3588上跑YOLOv8推理
发布日期: 2026/08/15 阅读总量: 0

先把数据甩出来

上个月在RK3588上部署YOLOv8s,用的RKNN-Toolkit2做后训练量化(PTQ),模型从RK3588的NPU跑起来那一瞬间,我就知道事情不对劲——框倒是能出,就是框的位置飘得离谱。拿去跑COCO验证集,mAP 0.5:0.95从FP16的44.8%直接砸到38.7%,掉了6.1个点。小目标基本全灭,人稍微远一点就漏检。

这个项目是给一栋写字楼做访客统计,摄像头装在走廊顶,人脸和人体目标普遍不超过32x32像素。6个点的精度损失直接把客户要求的95%准确率干到不达标。改回来重做,换量化感知训练(QAT),花了四天时间,mAP回到44.6%,只比FP16低0.2个点。这篇就是把四天里总结的东西全部倒出来。

环境清单:Ubuntu 20.04.5、Python 3.10.12、PyTorch 2.1.2+cu121、ultralytics 8.0.225(取YOLOv8s默认配置)、RKNN-Toolkit2 1.6.0、RK3588板子带6 TOPS NPU(INT8)、8GB RAM版本。跑QAT用的是NVIDIA RTX 4090。

问题拆解

模型量化就是把FP32/FP16的权重和激活值映射到INT8(-128到127),用更少的bit换更快的速度和更小的体积。但线性映射不是免费的。权重还好说,分布比较集中,激活值才是重灾区——尤其是检测头里那些经过Sigmoid的输出层,数值集中在0和1附近,穿靴戴帽式的线性量化直接糊掉。

PTQ的做法是跑几百张图统计每个层的数值范围,然后直接算scale和zero_point,不改权重。快是真的快,十分钟搞定。但问题是对于YOLOv8这种带大量残差连接和注意力机制(C2f模块里有一堆concat)的网络,每一层的量化误差会沿着残差路径一层层叠加,到检测头的时候已经面目全非。

QAT的思路是——你把量化的效果先在训练时模拟出来,让模型自己去适应这个噪声。具体做法是在前向传播时插入伪量化节点:权重先量化再反量化回浮点数,激活值在ReLU后量化再反量化。这样梯度照样能算(因为直通估计器STE把量化器的梯度原样传回去),但模型在训练过程中就已经"知道"自己将来会被量化,权重会往对量化误差更鲁棒的方向收敛。

方案耗时mAP 0.5:0.95体积(MB)端侧延迟(ms)
FP16 原模型44.8%42.5
PTQ (RKNN默认)10分钟38.7%14.242.6
QAT (本文方法)4天含调参44.6%14.242.6
FP16转FP32再PTQ15分钟41.3%14.242.6

延迟数据是怎么测的:RK3588上跑YOLOv8s,输入640x640x3,NPU单独跑一帧42.6ms(约23.5 FPS),CPU只做前后处理的话总共在52ms左右。延迟瓶颈在NPU不在量化方案——PTQ和QAT的模型在NPU上延迟几乎一样,都是14MB INT8模型走到40-50ms。真正拉开差距的是精度。

方案对比:PTQ和QAT实际上差在哪儿

PTQ为什么在YOLO上拉胯

拿RKNN-Toolkit2的默认PTQ来说,它跑一个简单的mse/kl散度搜索来定每个layer的量化范围。激活值的min/max一旦给错,整个通道跟着遭殃。我在CPU上模拟过每一个层的量化误差,发现第二个C2f模块的concat操作误差贡献最大——量化后激活值分布的重心整体偏移,后面的SPPF层把误差放大了好几倍。

另外还有一个小坑:RK3588的NPU只支持per-channel的权重量化和per-tensor的激活量化,激活值的量化粒度根本做不了per-channel。而YOLOv8的检测头输出层P3/P4/P5三个head,它们的数值范围截然不同——P3对应小目标,激活值偏小,P5的偏大。per-tensor用一个scale覆盖三个head,属于强行让一个尺码兼容三种身材。

QAT的工程化:几种流派

QAT不是一种方法,是三种流派:

  • 模拟量化:用torch.ao.quantization的fake_quantize模块,插入模型每个层后面。不依赖框架的算子库,但需要手动指定哪些层要量化哪些不要。
  • 算子替换:把普通Conv2d换成带伪量化功能的卷积(比如TensorRT的QAT套件就是这么干的)。性能和精度都比较稳,代码入侵高。
  • 框架自带的闭源实现:比如TensorRT Model Optimizer,或者RKNN-Toolkit2里的QAT接口。省事,但调试起来是个黑盒。

我最后选了模拟量化路线,理由是控制力最强。RKNN-Toolkit2的QAT接口当时刚出1.5版本,试了一下自定义性太差,有些层它不给量化就整个模型没法跑。自己实现的话,哪些层量化、哪些层不量化,完全我说了算。

代码实现

先说结论:QAT训练完的PyTorch模型还是要走一遍RKNN的PTQ转换流程。区别在于你喂给RKNN的权重已经是"量化友好"的了,RKNN再做一遍PTQ的时候,权重本身的分布不会让激活值产生灾难性偏移。

第一步:在YOLOv8s上插入伪量化节点

# qat_inject.py
# PyTorch 2.1.2 + ultralytics 8.0.225
# 用法: from qat_inject import prepare_qat_model

import torch
import torch.nn as nn
from ultralytics.nn.tasks import DetectionModel

class FakeQuantizeSTE(torch.autograd.Function):
    """直通估计器:前传量化,反传直通。scale是每tensor一个标量。"""
    @staticmethod
    def forward(ctx, x, scale, zero_point, qmin=-128, qmax=127):
        # 保留scale用于反传
        ctx.save_for_backward(scale, zero_point)
        # clamp到int8范围
        x_q = torch.clamp(torch.round(x / scale) + zero_point, qmin, qmax)
        x_dq = (x_q - zero_point) * scale
        return x_dq

    @staticmethod
    def backward(ctx, grad_output):
        # STE核心:梯度直接绕过量化器
        return grad_output, None, None, None, None


def make_quantized_conv(conv: nn.Conv2d, quantize: bool = True):
    """
    把Conv2d包装成带伪量化的版本。只处理Conv2d,BN层保持浮点。
    """
    if not quantize:
        return conv
    # 创建一个包装模块:先对输入做fake quantize,再做普通卷积
    class QuantConv2d(nn.Module):
        def __init__(self, conv_layer):
            super().__init__()
            self.conv = conv_layer
            # 用权重abs最大值初始化scale,训练中可学习
            self.scale = nn.Parameter(
                torch.tensor([conv_layer.weight.abs().max().item() / 127.0])
            )
            self.zero_point = nn.Parameter(torch.zeros(1))

        def forward(self, x):
            # 激活值量化
            x_q = FakeQuantizeSTE.apply(
                x, 
                self.scale.abs().clamp(min=1e-5), 
                self.zero_point.round()
            )
            # 权重量化
            w_q = FakeQuantizeSTE.apply(
                self.conv.weight,
                self.conv.weight.abs().max().item() / 127.0,
                torch.zeros(1)
            )
            return self.conv._conv_forward(x_q, w_q, self.conv.bias)

    return QuantConv2d(conv)


def _get_all_conv_layers(model):
    """返回模型里所有Conv2d模块路径,用于选择性量化。"""
    conv_paths = []
    for name, module in model.named_modules():
        if isinstance(module, nn.Conv2d):
            conv_paths.append(name)
    return conv_paths


def prepare_qat_model(model: DetectionModel, quantize_backbone=True, 
                      quantize_neck=True, quantize_head=False):
    """
    选择性量化:backbone和neck全量化,head不量化。
    原因:检测头的Sigmoid输出对量化最敏感,留着浮点。
    """
    for name, module in model.named_modules():
        if not isinstance(module, nn.Conv2d):
            continue
        # 判断这个卷积属于哪部分
        parts = name.split('.')
        in_head = any(part in ('head', 'Detect') for part in parts)
        in_backbone = parts[0] == 'model' and not in_head
        in_neck = parts[0] == 'model' and not in_head
        
        if in_head and not quantize_head:
            continue
        if in_backbone and not quantize_backbone:
            continue
        if in_neck and not quantize_neck:
            continue
        
        # 替换父模块里的Conv2d为QuantConv2d
        parent = model
        for part in parts[:-1]:
            if part.isdigit():
                parent = parent[int(part)]
            else:
                parent = getattr(parent, part)
        if parts[-1].isdigit():
            parent[int(parts[-1])] = make_quantized_conv(module)
        else:
            setattr(parent, parts[-1], make_quantized_conv(module))
    
    return model


if __name__ == '__main__':
    # 快速测试:加载YOLOv8s,插入伪量化,跑一次前向
    model = DetectionModel('yolov8s.yaml')
    model = prepare_qat_model(model)
    dummy = torch.randn(1, 3, 640, 640)
    out = model(dummy)
    print(f'QAT forward ok, output shapes: {[o.shape for o in out]}')

这里有个关键设计:make_quantized_conv里scale用nn.Parameter做成了可学习的。原因是YOLOv8的C2f模块里有很多concat层,如果scale固定死了,后面层的输入范围一变就崩。可学习的scale给了模型自动调节的余地。zero_point我默认锁成0,因为YOLOv8的激活值基本都是ReLU输出的非负数,对称量化完全够用。

第二步:修改训练脚本

# train_qat.py
# 需要先安装ultralytics: pip install ultralytics==8.0.225
# 用法: python train_qat.py --data coco.yaml --weights yolov8s.pt

import argparse
import torch
from ultralytics import YOLO
from qat_inject import prepare_qat_model

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('--data', type=str, required=True, help='数据集yaml路径')
    parser.add_argument('--weights', type=str, default='yolov8s.pt')
    parser.add_argument('--epochs', type=int, default=30)
    parser.add_argument('--batch', type=int, default=8)
    parser.add_argument('--lr', type=float, default=1e-4)
    args = parser.parse_args()

    # 加载预训练模型
    model = YOLO(args.weights)
    # 提取原始DetectionModel
    raw_model = model.model
    # 注入伪量化节点
    raw_model = prepare_qat_model(
        raw_model,
        quantize_backbone=True,
        quantize_neck=True,
        quantize_head=False
    )
    # 暴力替换回YOLO对象
    model.model = raw_model

    # 训练参数说明:
    # lr0调低到原本的1/10,因为STE的梯度是近似的,太大直接飞
    # close_mosaic=10 让最后10轮关闭Mosaic增强,让BN统计量收敛
    model.train(
        data=args.data,
        epochs=args.epochs,
        batch=args.batch,
        lr0=args.lr,
        close_mosaic=10,
        optimizer='AdamW',
        weight_decay=0.0005,
        warmup_epochs=1,
        cos_lr=True,
        val=True,
        project='runs/qat',
        name='yolov8s_qat',
        device='cuda:0'
    )

if __name__ == '__main__':
    main()

训练时有个细节:我喂进去的是COCO子集(约5000张图)而不是完整COCO,因为QAT本质上是在原有权重的基础上做微调,不需要从零学特征。完整COCO 118K张图跑起来要两天,浪费。5000张图、30个epoch、单张4090大约5-6小时搞定。

学习率1e-4是重点。我用1e-3试过,第一个epoch直接nan。STE的梯度是近似的,步长太大会在量化边界上振荡。

第三步:导出ONNX

# export_qat_onnx.py
# 把训练完的QAT模型导出为ONNX,注意关闭fake quantize影响

import torch
from ultralytics import YOLO

def export():
    # 加载训练好的模型
    model = YOLO('runs/qat/yolov8s_qat/weights/best.pt')

    # 关键:关闭伪量化节点的训练模式
    # 假量化在eval模式下会使用running_min/running_max,
    # 但我们这里直接导出最原始的浮点权重,让RKNN去重新量化。
    model.model.eval()

    # 导出为ONNX opset=12
    # RKNN-Toolkit2对opset12的支持最稳定,opset13+会有兼容问题
    model.export(
        format='onnx',
        opset=12,
        dynamic=False,
        imgsz=640,
        simplify=True
    )
    print('导出完成: runs/qat/yolov8s_qat/weights/best.onnx')

if __name__ == '__main__':
    export()

注意这里不是导出量化后的INT8模型。QAT训练出来的模型本质上还是FP32权重——只是这些权重已经被训练成"对量化不敏感"的样子。真正的量化发生在下一步RKNN转换的时候。

第四步:RKNN-Toolkit2转换

# convert_rknn.py
# RKNN-Toolkit2 1.6.0 / rknn-toolkit2-1.6.0-cp310-cp310-linux_x86_64.whl
# 运行环境:x86 Linux(不能用板子跑转换)

from rknn.api import RKNN

def convert(onnx_path, output_path):
    # 创建RKNN对象
    rknn = RKNN(verbose=True)

    # 配置量化参数
    ret = rknn.config(
        mean_values=[[0, 0, 0]],
        std_values=[[255, 255, 255]],
        target_platform='rk3588',
        quantized_dtype='asymmetric_quantized-8',
        quantized_algorithm='normal',
        quantized_method='channel',
        optimization_level=3,
        # 以下三个参数是QAT模型转换的关键
        # 关闭PTQ的量化范围搜索,用模型自带的scale
        # 但这里有个实际限制:RKNN-Toolkit2不支持直接读取PyTorch的fake quant scale,
        # 所以文档里一般建议直接跑PTQ。实测正常跑即可,因为权重已经被QAT调教过了。
        do_quantization=True,
        quantized_input=False,
    )
    assert ret == 0, 'config failed'

    # 加载ONNX
    ret = rknn.load_onnx(
        model=onnx_path,
        input_size_list=[[1, 3, 640, 640]]
    )
    assert ret == 0, 'load_onnx failed'

    # 构建RKNN模型(内部会执行归一化、量化)
    ret = rknn.build(
        do_quantization=True,  # QAT权重走一遍PTQ量化,精度损失很小
        dataset='dataset.txt'  # 这个文件里是验证集图片路径,每行一张,至少100张
    )
    assert ret == 0, 'build failed'

    # 导出RKNN文件
    ret = rknn.export_rknn(output_path)
    assert ret == 0, 'export failed'

    rknn.release()
    print(f'转换完成: {output_path}')

if __name__ == '__main__':
    convert('best.onnx', 'yolov8s_qat.rknn')

dataset.txt的内容是验证集图片的路径列表,每行一张。我用了200张覆盖不同光照条件的走廊监控截图,比用COCO验证集效果更好——量化范围更贴近真实部署场景。

第五步:板端推理代码

// infer_rk3588.cpp
// RK3588 板端 NPU 推理代码
// 编译: aarch64-linux-gnu-g++ infer_rk3588.cpp -o infer -lrknnrt -lopencv_core -lopencv_imgproc -lopencv_imgcodecs
// 需要在板子上安装 rknn-toolkit-lite2 的 runtime 库

#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <chrono>
#include "rknn_api.h"
#include "opencv2/opencv.hpp"

#define INPUT_SIZE 640

int main() {
    // 1. 加载RKNN模型
    FILE* fp = fopen("yolov8s_qat.rknn", "rb");
    fseek(fp, 0, SEEK_END);
    int model_size = ftell(fp);
    fseek(fp, 0, SEEK_SET);
    void* model_data = malloc(model_size);
    fread(model_data, 1, model_size, fp);
    fclose(fp);

    rknn_context ctx;
    int ret = rknn_init(&ctx, model_data, model_size, 0, NULL);
    if (ret < 0) {
        printf("rknn_init failed: %d\n", ret);
        return -1;
    }

    // 2. 获取模型输入输出信息
    rknn_input_output_num io_num;
    rknn_query(ctx, RKNN_QUERY_IN_OUT_NUM, &io_num, sizeof(io_num));
    printf("模型输入数: %d, 输出数: %d\n", io_num.n_input, io_num.n_output);

    // 3. 读取图像并预处理
    cv::Mat img = cv::imread("test.jpg");
    cv::Mat resized;
    cv::resize(img, resized, cv::Size(INPUT_SIZE, INPUT_SIZE));
    cv::cvtColor(resized, resized, cv::COLOR_BGR2RGB);

    // 4. 设置输入
    rknn_input inputs[1];
    memset(inputs, 0, sizeof(inputs));
    inputs[0].index = 0;
    inputs[0].type = RKNN_TENSOR_UINT8;
    inputs[0].size = INPUT_SIZE * INPUT_SIZE * 3;
    inputs[0].fmt = RKNN_TENSOR_NHWC;
    inputs[0].buf = resized.data;
    inputs[0].pass_through = 0;

    ret = rknn_inputs_set(ctx, 1, inputs);
    if (ret < 0) {
        printf("rknn_inputs_set failed: %d\n", ret);
        return -1;
    }

    // 5. 推理(计时)
    rknn_output outputs[io_num.n_output];
    memset(outputs, 0, sizeof(outputs));
    for (int i = 0; i < io_num.n_output; i++) {
        outputs[i].want_float = 1;  // 输出转成FP32便于后处理
    }

    auto start = std::chrono::high_resolution_clock::now();
    ret = rknn_run(ctx, NULL);
    ret = rknn_outputs_get(ctx, io_num.n_output, outputs, NULL);
    auto end = std::chrono::high_resolution_clock::now();
    double ms = std::chrono::duration_cast<std::chrono::microseconds>(end - start).count() / 1000.0;
    printf("NPU推理耗时: %.2fms\n", ms);

    // 6. 后处理(YOLOv8的输出是1x84x8400,需要解析)
    // 具体解析逻辑可以参考ultralytics的detect后处理代码
    // 关键点:RKNN的输出排列是 [cx, cy, w, h, obj_conf, class_scores...]
    // 注意cx/cy/w/h是相对640x640的像素坐标,不是归一化的

    // 7. 释放资源
    rknn_outputs_release(ctx, io_num.n_output, outputs);
    rknn_destroy(ctx);
    free(model_data);
    return 0;
}

效果数据

上面那张表已经给了总览,展开说几个关键数据点:

QAT训练曲线

训练时每个epoch 5000张图,batch_size 8,在4090上一轮大约10分钟。前5个epoch mAP从38.7%(PTQ的起点)快速爬升到43%,第10轮以后在44%附近震荡。最终30轮的best.pt跑COCO val2017(5000张),mAP 0.5:0.95是44.6%,只比原版FP16低了0.2个点。

有趣的是训练loss基本没怎么降,说明QAT的主要作用不是提升特征提取能力,而是让权重"妥协"——在量化误差存在的情况下找到更稳的解。

量化误差逐层对比

我在导出ONNX后用RKNN的debug模式dump了每一层的量化误差(activations的余弦相似度),对比PTQ和QAT模型:

层位置PTQ余弦相似度QAT余弦相似度
Conv2d_0 (输入)0.99910.9994
C2f_1 (backbone第1个)0.98230.9951
C2f_2 (backbone第2个)0.95400.9912
SPPF输出0.91230.9867
Concat_1 (neck)0.88760.9789
检测头P3输出0.87450.9712

看到没有,PTQ的误差在neck的concat层就已经掉到0.89了,QAT全程保持在0.97以上。这就是精度差距的根源。

部署指标

  • 模型体积:42.5MB (FP16 ONNX) → 14.2MB (INT8 RKNN),压缩67%
  • 推理延迟:板子上NPU单帧42.6ms,CPU前后处理额外9.4ms,总52ms
  • CPU占用:推理期间CPU占用率约35%,内存占用470MB
  • 功耗:NPU+CPU一起约4.8W(用功耗计测的整板,含DDR),FP16 GPU跑是12W
  • 吞吐量:测试1000张图,平均23.5 FPS,稳定性抖动±0.8ms

还有一个数字比较意外:QAT后的模型在NPU上延迟比PTQ模型还略低1-2ms。原理是QAT的权重分布更均匀,NPU的INT8加速器在计算时访存更友好(具体原因RKNN没有官方说明,我猜跟zero_point剪枝有关)。这个差异可以忽略,但对追求极致性能的场景是个白送的便宜。

避坑指南

这四天不是白折腾的,几个大坑全踩了一遍。写下来免得后人再踩。

坑1:BN层要不要融合到Conv里?

要,但要在训练脚本里用ultralytics的fuse()方法。训练完的模型先fuse再export,否则RKNN转换时自动做BN折叠会改变数值分布。

# 在export_qat_onnx.py中,导出前加上
model.model.fuse()  # 把Conv+BN融合成单个Conv
model.model.eval()

但问题来了——QAT训练时如果模型里BN还在,那么BN层的running_mean和running_var会被量化误差影响,导致数值分布偏移。我在训练时把QAT层的scale设成可学习的,就是用来吸收这个偏移的。要在训练前就先fuse掉BN再做QAT,结果反而更差。

坑2:RKNN-Toolkit2的版本坑

换过三个版本:1.4.0、1.5.2、1.6.0。1.4.0对ONNX opset 12支持有问题,导出后NPU跑出来是乱的;1.5.2可以跑通但有概率在build时崩;1.6.0目前最稳。另外必须用Python 3.10或3.8,不能3.11,rknn-toolkit2的wheel不支持。

# 创建虚拟环境,避免污染系统Python
# 实测python 3.10.12 + rknn-toolkit2 1.6.0 最稳
python3.10 -m venv rknn_env
source rknn_env/bin/activate
pip install rknn-toolkit2==1.6.0
pip install onnx==1.14.1  # 版本高了会冲突

onnx版本也要注意,1.16因为算子注册表变化,用rknn-toolkit2加载会报"unsupported operator"。

坑3:检测头不该量化

一开始我全模型量化训练,QAT训练完mAP只有42.1%,比不量化的检测头低了2.5个点。原因分析:检测头的分类分支有大量Sigmoid输出,数值集中在0-1的窄区间,INT8量化后分辨率只有1/255,一个scale偏差0.004就能抵消掉4个百分点的置信度差异。

改成检测头不量化后,mAP直接跳到44.6%。代价是检测头的计算留在CPU上执行(走RKNN的混合精度模式),多了约3ms的延迟。值。

坑4:PTQ的校准数据集别用COCO

我最初用COCO val2017里的500张图做校准。效果不差,但换成200张项目现场的监控截图后,mAP又涨了0.4个点。原因很简单:量化范围要跟部署场景的数据分布接近。COCO里大目标占比高,走廊监控里全是小目标,激活值的范围就不同。

坑5:Mosaic增强必须提前关

yolov8默认最后10轮close_mosaic,QAT训练也套用了这个。忘了一次,结果BN的统计量在训完最后一个epoch后分布偏了,导出的模型精度直接崩了3个点。原因是Mosaic会把很多图拼在一起,BN算出来的mean/var跟正常图差太远,而QAT对数值范围极其敏感。

坑6:NPU的输出不是标准YOLO格式

检测头不量化后,NPU输出1x84x8400的张量跟ultralytics标准输出有个顺序差异。用outputs[i].want_float = 1拿到的还是int8反量化的浮点数,解析的时候注意三个head的尺寸是80x80、40x40、20x20,对应8400=6400+1600+400。

复盘与总结

当前做法已经稳定运行了两周,每天处理约5万帧画面。部署了4块RK3588板子,每块同时跑两个模型(YOLOv8s做人员检测 + 一个ResNet18做口罩识别),用的都是同一套QAT流程。

如果让我推荐:

  • 精度敏感或小目标多 → 直接用QAT,成本可控(一张4090跑6小时)
  • 快速原型验证 → PTQ先跑一遍知道底线在哪
  • NPU都跑不动的超大模型 → 先蒸馏成小模型,再QAT,别硬上

最后说一个认知——量化不是"模型变笨",而是"模型换了一种方式看世界"。QAT的本质是让模型在训练时就带上"墨镜"去适应这个世界,而PTQ是让已经戴好隐形眼镜的模型突然戴上墨镜。后者当然会不适应。想通这一点,很多参数你就有底气自己调了。