先把数据甩出来
上个月在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.2 | 42.6 |
| QAT (本文方法) | 4天含调参 | 44.6% | 14.2 | 42.6 |
| FP16转FP32再PTQ | 15分钟 | 41.3% | 14.2 | 42.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.9991 | 0.9994 |
| C2f_1 (backbone第1个) | 0.9823 | 0.9951 |
| C2f_2 (backbone第2个) | 0.9540 | 0.9912 |
| SPPF输出 | 0.9123 | 0.9867 |
| Concat_1 (neck) | 0.8876 | 0.9789 |
| 检测头P3输出 | 0.8745 | 0.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是让已经戴好隐形眼镜的模型突然戴上墨镜。后者当然会不适应。想通这一点,很多参数你就有底气自己调了。