本文完整的记录了基于Ultralytics项目中的 RT-DETR 预训练模型在自定义数据集上进行的训练、ONNX导出,以及在RK3588硬件平台上的适配流程。
对于 RT-DETR 模型的实现而言,目前主要有两套独立的实现,即 Baidu 的 PaddleDetection 和 Ultralytics (YOLOv8/v11 生态)。因为我自己对于 Ultralytics/YOLO 较为熟悉,所以本文以 Ultralytics/YOLO 的实现进行试验和记录。

因为这里使用的是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没有多大区别,有两点注意事项:
fork,它会直接克隆当前的内存状态,不需要重新导入模块;而 Windows 不支持 fork,只能 spawn(冷启动一个新的 Python 解释器并导入以上脚本),这样的话,子进程启动时会再次运行一遍整个 train.py,从而又触发了一个新的训练请求,形成死循环。系统为了保护内存不被耗尽,主动抛出 RuntimeError来结束整个程序。而把训练代码包裹在__main__里面以后,只有主进程会执行train_model函数,子进程只负责执行训练过程,这样就不会进入反复新建子进程的死循环。以上的训练执行过程明显比使用相同数据集训练的YOLOv8模型耗时要长一些,应该是RT-DETR模型的计算复杂度要高于 YOLOv8。训练结果同样保存在runs目录下的best.pt文件。
在RK3588硬件上跑AI模型,大致的路径无非是:
rknn-toolkit-lite2 工具(Python)或者C/C++语言把模型跑起来,在板端完成推理测试得到结果。所以既然前面的训练过程已经完成,那么下一步自然就是进行ONNX格式的模型转换。但是这一步出现了一个由Grid Sample算子所导致的很大的问题: