利用ultralytics版本的RT-DERT训练NEU-DET缺陷数据集（windows版本和服务器版本教程）

数据集

NEU-DET缺陷数据集是东北大学（NEU）发布的表面缺陷数据库，收集了热轧钢带的六种典型表面缺陷，即轧制氧化皮（RS），斑块（Pa），开裂（Cr），点蚀表面（ PS），内含物（In）和划痕（Sc）

地址1（源文件voc.XML格式）

链接: https://pan.baidu.com/s/1vRB0USgHY0n3tzW3CQOG9w?pwd=bnm2 提取码: bnm2

地址二（coco.json格式）

链接: https://pan.baidu.com/s/1FAjAZ76OumhCJvmN0WvkDQ?pwd=1ce3 提取码: 1ce3

地址三（yolo格式.txt格式）

链接: https://pan.baidu.com/s/1LgNKunzrWVcYLPxFL57pYw?pwd=cnun 提取码: cnun

三个地址的训练集，验证集和测试集已经划分完毕，直接可以使用。

RT-DERT模型版本

RT-DERT论文地址：https://arxiv.org/pdf/2304.08069

论文中的代码地址：https://github.com/lyuwenyu/RT-DETR

ultralytics版本地址：GitHub - ultralytics/ultralytics: NEW - YOLOv8 🚀 in PyTorch > ONNX > OpenVINO > CoreML > TFLite

为什么不用原论文给的代码训练而用内嵌在ultralytics代码训练？

论文代码目录

论文代码给了俩个框架训练，一个pdddle框架和torch训练，环境安装也很简单，但是太过单一，训练结束后没有显示进行预测的结果，只有损失值和.pth文件，不利于直观的表达，也没有图像化结果，我想大家还是希望有一个图形话结果去显示的。以下是pytorch训练结果的内容.

log.txt内容

论文代码的数据格式是JSON数据格式，也就是地址二的数据集，不要弄错了o

ultralytics版本代码目录

因为是YOLO8内嵌的RT-DETR,因此也更加完善，评价体系也更加全面，魔改也比较方便，因此大多数博主推荐的用ultralytics版本进行RT-detr的训练，那就圆规正转，接下来我将详细说明如何在windows上训练模型。

训练模型

第一步：先下好压缩包

点击ultralytics版本地址然后点code下载压缩包。

然后解压文件夹

第二步：放置数据集和新建代码

ultralytics版本支持yolo格式的数据集，因此大家需要把地址三的数据集下载好，如果会格式转换的话就不用这么麻烦了。然后新建文件夹。

data文件夹，NEU-DET.yaml和train.py，其他多出来的为训练过程中下载的东西

data文件夹目录

NEU-DET.yaml，训练路径配置，NEU-DET.yaml放在\ultralytics-main\ultralytics\cfg\datasets\NEU-DET.yamll

path: C:/Users/PMgood/Desktop/ultralytics-main/data/NEU-DET  # dataset root dir
train: C:/Users/PMgood/Desktop/ultralytics-main/data/NEU-DET/train/images  # train labels
val: C:/Users/PMgood/Desktop/ultralytics-main/data/NEU-DET/val/images  # val labels
#改成你自己的路径

# number of classes
nc: 6

# class names
names:
  0: crazing
  1: inclusion
  2: patches
  3: pitted_surface
  4: rolled-in_scale
  5: scratches

train.py代码

from ultralytics.cfg import entrypoint

arg = "yolo detect train model=rtdetr-l.yaml data=ultralytics/cfg/datasets/NEU-DET.yaml"

entrypoint(arg)

配置环境

环境配置的pytorch环境

我的cuda支持最高版本为12.3，可以打开命令框输入nvidia-smi查看。因此我支持的cuda版本为<=12.3,所以我安装了12.1版本的，python版本为3.80.这里我就不多阐述怎么安装pytorch环境。需要的可以参考https://zhuanlan.zhihu.com/p/612181449，很实用。

pytorch版本为

pip install torch==2.1.0 torchvision==0.16.0 torchaudio==2.1.0 --index-url https://download.pytorch.org/whl/cu121

安装ultralytics

pip install ultralytics

ultralytics库是一个集成库，包含了很多的库，因此就不需要一个一个安装了

现在我们已经配置好环境了，可以检查自己的环境是否配好。

开始训练

运行我们新建的train.py，文件显示出错

显示找不到文件，因此我们需要打开ultralytics/settings.yaml,文件，把datasets_dir路劲改成和NEU-DET.yaml path路径一样

path: C:/Users/PMgood/Desktop/ultralytics-main/data/NEU-DET

修改之后

接着训练显示出错

因为window不能多线程，因此需要把worker设置为0,原来是8

然后继续运行，显示正在训练就成功了。如果内存溢出的话减小bath size和图像分辨率，但是太慢了，因此可以搬到云服务器上运行

云服务器上训练

过程还是不变的，也需要更改数据集路径和ultralytics/settings.yaml路径，之后就可以直接运行了。

运行结果

然后就可以下载runs文件进行分析了。