RKNN3588上部署 RTDETRV2

Mrs.Gril

于 2025-06-05 16:04:26 发布

阅读量1.2k

点赞数 13

CC 4.0 BY-SA版权

分类专栏： AI-训练+部署文章标签： YOLO RTDETR 深度学习 RKNN

本文链接：https://blog.youkuaiyun.com/zhangdaoliang1/article/details/148452884

AI-训练+部署专栏收录该内容

58 篇文章

订阅专栏

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 sys

import torch
import torch.nn as nn

from src.core import YAMLConfig
from rtModel import Model

def 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 onnx
        onnx_model = onnx.load(output_file)
        onnx.checker.check_model(onnx_model)
        print('Check export onnx model done...')

    if args.simplify:
        import onnx
        import onnxsim
        dynamic = True
        # input_shapes = {'images': [1, 3, 640, 640], 'orig_target_sizes': [1, 2]} if dynamic else None
        input_shapes = {'images': data.shape, 'orig_target_sizes': size.shape} if dynamic else None
        onnx_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 argparse
    parser = 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

得到结果：