数据集
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文件进行分析了。