RKNN3588上部署 RTDETRV2
RT-DETR V2 是由百度研究团队在 2024年 提出的,是其广受好评的实时目标检测模型 RT-DETR 的重大升级版本。它继承了第一代 RT-DETR 利用 Transformer 架构实现端到端目标检测 和 卓越实时性能 的核心优势,并针对模型精度、训练效率和部署灵活性进行了全方位的显著提升,再次刷新了实时检测的SOTA(State-of-the-Art)水平。
核心目标与核心优势:
-
更高精度: 在保持甚至超越实时速度的前提下,显著提升在各种场景下的检测精度(平均精度均值 - AP)。
-
更快训练: 极大地缩短模型训练所需时间,降低研发和迭代成本。官方宣称训练速度比 V1 快 3倍。
-
更强部署灵活性: 提供更广泛的硬件平台支持(如 NVIDIA GPU, Intel CPU, NVIDIA Jetson, ARM 处理器等)和更高效的推理引擎适配(如 TensorRT, ONNX Runtime, OpenVINO)。
关键技术革新:
RT-DETR V2 的成功源于多项创新设计:
-
混合编码器设计:
-
巧妙地融合了 CNN 骨干网络(如 HGNetv2) 和 Transformer 编码器 的优势。
-
CNN 高效提取局部特征,Transformer 则擅长捕捉长距离依赖和全局上下文。
-
这种混合结构在特征提取阶段就为高精度和高效性奠定了基础。
-
-
不确定性最小化查询(Uncertainty-minimal Queries - UMQs):
-
这是 V2 的核心创新之一。传统 DETR 类模型使用可学习的对象查询(Object Queries)来定位和识别目标,但其初始化和匹配存在不确定性。
-
UMQs 机制通过 显式地建模和最小化查询的空间位置不确定性,引导模型学习到更稳定、更精准的初始目标位置表示。
-
效果: 显著加速了训练收敛速度(是训练提速的关键),并提升了检测精度,尤其是对小目标的检测能力。
-
-
IoU 感知查询选择:
-
在将特征图上的点(Anchor Points)转换为解码器输入的对象查询时,V2 引入了 IoU(交并比)感知机制。
-
它不仅考虑分类得分,还综合考虑预测框与真实框的IoU质量来选择最具有代表性和高质量的查询。
-
效果: 筛选出更优质的初始查询,为解码器提供更有效的输入,进一步提升精度。
-
-
尺度增强的混合分配:
-
改进了训练过程中为预测框分配监督标签(即正负样本分配)的策略。
-
通过结合动态匹配(如匈牙利算法)和基于规则的静态分配(如ATSS)的优点,并融入多尺度信息。
-
效果: 为模型提供了更丰富、更稳定的监督信号,尤其是在处理不同尺度的目标时,提升了训练的鲁棒性和最终精度。
-
性能表现:
在权威的 COCO 目标检测数据集上,RT-DETR V2 展现了令人印象深刻的性能:
-
例如,RT-DETR-L 模型在 COCO val2017 上达到 54.8% AP,同时推理速度在 NVIDIA Tesla T4 GPU 上高达 74 FPS。
-
更重要的是,它在 各种硬件平台(从高端GPU到边缘CPU) 上都实现了领先的精度-速度权衡,验证了其卓越的部署灵活性。
-
其训练速度相比 V1 大幅提升(约3倍),极大地提高了研发效率。
但是在边缘端上部署确实存在较多的问题,今天我们在RK3588上部署RTDETRV2.
导出onnx:
import os
import sysimport torch
import torch.nn as nnfrom src.core import YAMLConfig
from rtModel import Modeldef main(args, ):"""main"""cfg = YAMLConfig(args.config, resume=args.resume)model = Model()data = torch.rand(1, 3, 1024, 1024)size = torch.tensor([[1024, 1024]])_ = model(data, size)dynamic_axes = {'images': {0: 'N', },'orig_target_sizes': {0: 'N'}}output_file = args.resume.replace('.pth', '.onnx') if args.resume else 'model.onnx'torch.onnx.export(model,(data, size),output_file,input_names=['images', 'orig_target_sizes'],output_names=['labels', 'boxes', 'scores'],# dynamic_axes=dynamic_axes,opset_version=16,verbose=False,do_constant_folding=True,)if args.check:import onnxonnx_model = onnx.load(output_file)onnx.checker.check_model(onnx_model)print('Check export onnx model done...')if args.simplify:import onnximport onnxsimdynamic = True# input_shapes = {'images': [1, 3, 640, 640], 'orig_target_sizes': [1, 2]} if dynamic else Noneinput_shapes = {'images': data.shape, 'orig_target_sizes': size.shape} if dynamic else Noneonnx_model_simplify, check = onnxsim.simplify(output_file, test_input_shapes=input_shapes)onnx.save(onnx_model_simplify, output_file)print(f'Simplify onnx model {check}...')if __name__ == '__main__':import argparseparser = argparse.ArgumentParser()parser.add_argument('--config', '-c', default='configs/dfine/dfine_hgnetv2_l_coco.yml', type=str, )parser.add_argument('--resume', '-r', type=str, )parser.add_argument('--check', action='store_true', default=True,)parser.add_argument('--simplify', action='store_true', default=True,)args = parser.parse_args()main(args)
转换为rknn,正常转换即可。
得到rknn将其push到开发板中。
执行命令:
./rknn_yolov5_demo ../REDETRV2_RKNN.rknn ../vedio_0_2_580.jpg 10 RTDETRV2
得到结果: