【代码复现】RT-DETR2（pytorch版）环境部署及训练

R_grace

已于 2025-04-01 16:42:04 修改

阅读量3.3k

点赞数 18

文章标签： python 深度学习

于 2024-12-05 23:02:20 首次发布

本文链接：https://blog.csdn.net/R_grace/article/details/144270083

版权

简介：RT-DETR是第一个实时端到端目标检测器，RT-DETR2 是基于 Transformer 的端到端检测器的 实时目标检测模型 (Real-Time Detection Transformer) ，目标是进一步提升目标检测的性能和效率，尤其是在实时目标检测场景下，相较于原始版本进行了性能和架构优化，其设计理念是以 高效性 和 实时性 为核心，结合 Transformer 模型的全局特征提取能力，为目标检测任务提供了兼具高性能和高速度的解决方案。

本教程配置：Windows10，CUDA11.8，python 3.8

1.代码拉取

代码链接：lyuwenyu/RT-DETR: [CVPR 2024] Official RT-DETR (RTDETR paddle pytorch), Real-Time DEtection TRansformer, DETRs Beat YOLOs on Real-time Object Detection. 🔥 🔥 🔥

git clone https://github.com/lyuwenyu/RT-DETR.git

2.虚拟环境

conda create -n rtv2 python=3.8 -y
conda activate rtv2

3.安装依赖

cd rtdetrv2_pytorch
pip install six
pip install opencv-python
pip install torch==2.1.0+cu118 torchvision==0.16.0+cu118 torchaudio==2.1.0+cu118 -f https://download.pytorch.org/whl/cu118/torch_stable.html
pip install -r requirements.txt
pip install scipy

检查CUDA是否可用的命令（可不执行）

import torch
print("Is CUDA available:", torch.cuda.is_available())  # 应返回 True
print("CUDA version:", torch.version.cuda)  # 查看 CUDA 版本
print("Device count:", torch.cuda.device_count())  # 查看可用 GPU 数量

4.制作数据集

1）coco数据集格式

位置：rtdetrv2_pytorch > configs > dataset

修改：yolo的label（txt）转coco的annotations（json）代码为yoyo2coco.py【网上随便都能搜到】

2）配置yaml

位置：rtdetrv2_pytorch > configs > dataset > coco_detection.yml

修改：（1）num_classes；（2）train_dataloader的img_folder和ann_file；（3）val_dataloader的img_folder和ann_file；

img_folder: ./configs/dataset/coco/train/
ann_file: ./configs/dataset/coco/annotations/instances_train.json

5.训练

1）修改train.py

位置：rtdetrv2_pytorch > tools > train.py

修改：第45行添加：default='./configs/rtdetrv2/rtdetrv2_r50vd_6x_coco.yml'

2）修改训练参数

位置：rtdetrv2_pytorch > configs > rtdetrv2 > include > dataloader.yml

修改：epoch

3）训练运行

位置：RT-DETR-main > rtdetrv2_pytorch > python tools/train.py

python tools/train.py