RT-DETR rtdetr-r18 - ultralytics - YOLOv8版本训练自己的数据集

CSPhD-winston-杨帆

已于 2025-04-20 21:35:31 修改

阅读量8.1k

点赞数 16

分类专栏：计算机视觉实践文章标签： YOLO

于 2024-01-05 11:04:32 首次发布

本文链接：https://blog.csdn.net/WhiffeYF/article/details/135371679

版权

计算机视觉实践专栏收录该内容

55 篇文章

订阅专栏

本文介绍了如何在Python环境中安装和配置RT-DETR模型，包括YOLOv8环境设置、数据集准备、以及如何使用私有GitHub仓库进行训练。同时，也提及了如何升级PyTorch到1.9.0版本以支持模型运行。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

0 源码与相关参考

源码请从b站up主【魔傀面具】获取：https://space.bilibili.com/286900343?spm_id_from=333.337.0.0

以下代码为私有

github：

https://github.com/Whiffe/ultralytics-RT-DETR

码云

https://github.com/Whiffe/ultralytics-RT-DETR

参考：b站魔傀面具
https://space.bilibili.com/286900343?spm_id_from=333.337.0.0

1 安装

注意，开始之前，需要配置好YOLOv8的环境

conda create -n rtdetr python=3.8

conda activate rtdetr

git clone https://github.com/Whiffe/ultralytics-RT-DETR
# git clone https://github.com/Whiffe/ultralytics-RT-DETR

cd ultralytics-RT-DETR

pip install ultralytics

python setup.py develop

pip install timm thop efficientnet_pytorch einops grad-cam dill -i https://pypi.tuna.tsinghua.edu.cn/simple

pip install -U openmim

mim install mmengine

mim install "mmcv>=2.0.0"

pip install psutil

2 数据集准备

moon.yaml

path: D:\mySearch\RT-detr\ultralytics-RT-DETR\dataset
train: ./yolo_behavior_Dataset_all2/images/train
val: ./yolo_behavior_Dataset_all2/images/val

# number of classes
nc: 5

# class names
names: [ 'weigh','height measure','drop ball','size measure','record']

3 训练 rtdetr-r18

python train.py

在这里插入图片描述

import warnings
warnings.filterwarnings('ignore')
from ultralytics import RTDETR

if __name__ == '__main__':
    model = RTDETR('ultralytics/cfg/models/rt-detr/rtdetr-r18.yaml')
    # model.load('') # loading pretrain weights
    model.train(data='dataset/moon.yaml',
                cache=False,
                imgsz=640,
                epochs=100,
                batch=4,
                workers=4,
                device='0',
                # resume='', # last.pt path
                project='runs/train',
                name='exp',
                )

在这里插入图片描述

4 pytorch 版本升级

pytorch离线安装网站：

https://download.pytorch.org/whl/torch_stable.html

pytorch需要1.9.0及以上版本

cu111/torch-1.9.0%2Bcu111-cp38-cp38-linux_x86_64.whl
https://download.pytorch.org/whl/cu111/torch-1.9.0%2Bcu111-cp38-cp38-linux_x86_64.whl

在这里插入图片描述
卸载

pip uninstall torch

安装

pip install torch-1.9.0+cu111-cp38-cp38-linux_x86_64.whl