RF-DETR自定义数据集并训练(实例测试）

原创

已于 2025-09-06 00:15:54 修改 · 510 阅读

6 ·

CC 4.0 BY-SA版权

文章标签：

#人工智能 #计算机视觉

于 2025-09-06 00:04:52 首次发布

RF-DETR支持多种流行的数据集格式，便于使用现有数据集或准备自定义数据。该框架设计用于处理COCO风格的标注，这已成为目标检测任务的事实标准。

支持的数据集格式

RF-DETR原生支持三种主要数据集格式：

COCO格式：目标检测数据集最广泛使用的格式
Roboflow格式：适用于Roboflow数据集的简化COCO兼容格式
Objects365格式：用于包含365个物体类别的大规模目标检测

本文采用Roboflow格式

fire_smog/
├── train/
│   ├── image1.jpg
│   ├── image2.jpg
│   ├── _annotations.coco.json
│   └── ...
├── valid/
│   ├── image1.jpg
│   ├── image2.jpg
│   ├── _annotations.coco.json
│   └── ...
└── test/
    ├── image1.jpg
    ├── image2.jpg
    ├── _annotations.coco.json
    └── ...

训练前需要安装环境pip install rfdetr，需要科学上网，我们这里用到了RFDETRSmall模型，可提前准备，没有的话，运行下面的训练脚本会自动下载

1.单卡训练方案单张3090

import os
# 只使用 GPU 1（或 GPU 2）
os.environ['CUDA_VISIBLE_DEVICES'] = '1'  
os.environ['PYTORCH_CUDA_ALLOC_CONF'] = 'expandable_segments:True'

import torch
torch.cuda