本文完整的记录了基于Ultralytics项目中的 RT-DETR 预训练模型在自定义数据集上进行的训练、ONNX导出,以及在RK3588硬件平台上的适配流程。

对于 RT-DETR 模型的实现而言,目前主要有两套独立的实现,即 Baidu 的 PaddleDetection 和 Ultralytics (YOLOv8/v11 生态)。因为我自己对于 Ultralytics/YOLO 较为熟悉,所以本文以 Ultralytics/YOLO 的实现进行试验和记录。

image.png

1.自定义数据集的训练

因为这里使用的是Ultralytics项目下的RT-DETR实现,所以其训练过程、训练所生成的中间文件组织结构以及对训练数据的表格格式、目录结构等都与YOLO模型没有任何区别,只需要把之前在YOLO模型上训练所使用的数据集目录以及其 data.yaml 直接拷贝过来使用即可。所以如果对YOLO模型的训练流程有经验的话,在这个项目中训练RT-DETR就很容易上手了。可参考一文总结YOLOv5模型的自定义训练、测试与模型导出的全过程。

训练流程的代码也与Ultralytics项目下的其他YOLO模型类似,训练过程的代码如下:

from ultralytics import RTDETR

def train_model():
    # 1. 加载预训练模型
    model = RTDETR('_my_research/pre_trained/rtdetr-l.pt') 

    # 2. 开始训练
    model.train(
        data='_my_research/rt_detr/data.yaml',
        epochs=100,
        imgsz=640,
        batch=4,
        lr0=0.0001,
        optimizer='AdamW',
        device=0,
        workers=2,  
        project='_my_research/runs',  # 指定输出目录
        name='rt_detr_train'  # 训练任务名称
    )

if __name__ == '__main__':
    train_model()

训练流程基本上与普通的YOLO没有多大区别,有两点注意事项:

以上的训练执行过程明显比使用相同数据集训练的YOLOv8模型耗时要长一些,应该是RT-DETR模型的计算复杂度要高于 YOLOv8。训练结果同样保存在runs目录下的best.pt文件。

2.Grid Sample算子的梦魇与Monkey Patch

在RK3588硬件上跑AI模型,大致的路径无非是:

所以既然前面的训练过程已经完成,那么下一步自然就是进行ONNX格式的模型转换。但是这一步出现了一个由Grid Sample算子所导致的很大的问题: