官方pytorch版本RT_DETR训练自己的数据集过程并导出结果

TracyGC

已于 2023-12-08 17:54:00 修改

阅读量2k

点赞数 10

分类专栏：研究与思考文章标签： pytorch 人工智能 python

于 2023-12-08 11:30:37 首次发布

本文链接：https://blog.csdn.net/tracygc/article/details/134872700

版权

37 篇文章 2 订阅

订阅专栏

本文介绍了如何使用RT-DETR在自定义数据集上进行训练，包括配置文件的修改、训练过程以及遇到的`onnx`模块错误的解决，最后详细说明了如何导出模型为ONNX格式。

摘要由CSDN通过智能技术生成

在前面的步骤中，利用配置好的COCO的数据集在“rtdetr_pytorch/configs/dataset/coco_detection.yml"里面把文件地址修改成自己的数据集地址，配制成my_coco.yml,

同时修改需要训练"rtdetr_pytorch/configs/rtder/rtdetr_r50vd_m_6x_coco.yml"里面的数据集

然后利用"RT-DETR-main/rtdetr_pytorch/tools/train.py，记得加入configs

训练结果会保存在以下的文件的地址(后面我们导出需要这个文件)

或者也可以直接利用edit configurations利用下面的代码运行：不过要注意加上eval参数

 -c configs/rtdetr/rtdetr_r101vd_6x_coco.yml -t model_pth/rtdetr_r101vd_2x_coco_objects365_from_paddle.pth

官方版本目前只支持导出ONNX格式，我们找到如下文件“RT-DETR-main/rtdetr_pytorch/tools/export_onnx.py”，将configs导入当时配置的文件，将训练完的权重文件输入到resume里，然后运行整个文件就会在同级目录下导出onnx文件如下：

出现了没有找到onnx:ModuleNotFoundError: No module named 'onnx'

pip install onnx==1.13.0
pip install paddle2onnx==1.0.5

再运行，就导出了model.onnx:

关注