OpenCV与AI深度学习 | 如何在自定义数据集上训练 RF-DETR

本文来源公众号“OpenCV与AI深度学习”，仅用于学术分享，侵权删，干货满满。

RF-DETR于 2025 年 3 月 19 日发布，是 Roboflow 开发的基于 Transformer 的对象检测模型架构。

RF-DETR 实现了最先进的性能，在 COCO 和新推出的 RF100-VL 数据集上均击败了 LW-DETR 和 YOLOv11 等模型。RF100-VL 是一个基准，旨在验证具有检测能力的模型在一系列领域的通用性。

官方代码链接：roboflow/rf-detr: RF-DETR is a real-time object detection model architecture developed by Roboflow, SOTA on COCO & designed for fine-tuning.

通过将分辨率扩展到 728，RF-DETR 在 NVIDIA T4 GPU 上以 25 FPS 实现了 60.5 mAP，成为第一个在 Microsoft COCO 基准上突破 60mAP 障碍的有记录的实时模型。该模型在 NVIDIA T4 上也实现了 25 FPS。

RF-DETR 是根据 Apache 2.0 许可证授权的，允许免费商业使用。

本文中我们将介绍如何在自定义数据集上训练 RF-DETR 模型。我们将以麻将牌识别模型为例进行训练，该任务涉及识别多个不同的类别。

以下是我们将要训练的模型的检测结果（右），以及标注结果（左）：

我们的模型结果几乎与事实相同，证明了 RF-DETR 预测的质量。

准备数据集

首先，我们需要准备一个数据集。在本指南中，我们将使用麻将牌识别数据集，这是RF100-VL 基准中的数据集之一。该数据集包含 2,000 多张麻将牌图像，并根据 Apache 2.0 许可证获得许可。

https://universe.roboflow.com/rf-100-vl/mahjong-vtacs-mexax-m4vyu-sjtd?ref=blog.roboflow.com

您可以从 Roboflow Universe 上的麻将牌页面下载数据集。在下面的训练部分中，我们将展示如何将数据集直接下载到 Colab 笔记本中。

您可以使用任何您想要的数据集来遵循本指南，但您需要一个COCO JSON 格式的数据集来训练 RF-DETR 模型。

如果您需要标记数据集，则可以使用Roboflow Annotate（我们功能齐全的基于 Web 的注释工具）进行标记。Annotate 附带一系列工具来加快标记过程，包括一套 AI 辅助标记工具。

如果您需要将数据集转换为 COCO JSON，您可以将其上传到 Roboflow 并以 COCO JSON 格式导出数据。

在设备上训练

我们准备了一个Colab 笔记本，您可以使用它来遵循本指南。我们的 Colab 笔记本介绍了从使用基本 COCO 权重尝试 RF-DETR 到在经过微调的模型上进行训练和运行推理的所有内容。

我们建议使用 NVIDIA A100 GPU 来微调 RF-DETR。

步骤1：安装RF-DETR SDK

首先，我们需要安装 RF-DETR SDK（python >=3.9即可）。您可以使用以下命令安装 SDK：

!pip install -q rfdetr

此外，运行 nvidia-smi 以确保您有可用的 GPU。输出应显示您的 GPU，如下所示：

步骤2：下载数据集

接下来，我们需要将数据集下载到我们的训练环境中。我们可以直接从 Roboflow 下载麻将数据集。

首先，在名为 ROBOFLOW_API_KEY 的环境变量中设置您的 Roboflow API 密钥。了解如何检索您的 Roboflow API 密钥。

然后运行以下代码：

from roboflow import download_datasetdataset = download_dataset("https://universe.roboflow.com/rf-100-vl/mahjong-vtacs-mexax-m4vyu-sjtd/dataset/2", "coco")

您可以将上面的数据集 URL 替换为 Roboflow Universe 上的 500,000 多个数据集中的任何一个。

步骤3：启动RF-DETR训练作业

在我们的训练环境中有了标记的数据集后，我们现在可以开始微调模型。

为了微调模型，我们可以使用以下代码：

from rfdetr import RFDETRBasemodel = RFDETRBase()history = []def callback2(data):	history.append(data)model.callbacks["on_fit_epoch_end"].append(callback2)model.train(dataset_dir=dataset.location, epochs=15, batch_size=16, lr=1e-4)

在这里，我们加载 RF-DETR 基础模型，然后传入我们下载的数据集的位置。在本指南中，我们将训练 15 个时期并使用 16 的批处理大小。此批处理大小针对 A100 进行了优化。

运行代码时，您将看到显示模型训练进度的消息，如下所示：

训练模型所需的时间取决于数据集的大小和您指定的周期数。在我们的测试中，使用 A100 GPU 训练一个包含 2,000 张图像的模型 15 个周期大约需要一个小时。

我们建议对生产模型进行至少 50 个时期的训练。

一旦您的模型训练完毕，模型权重和相关元数据将保存在名为 output 的目录中。

步骤4：审查模型评估指标

您可以通过绘制 RF-DETR 训练程序保存的数据来查看模型指标。

您可以使用以下代码绘制损失：

import matplotlib.pyplot as pltimport pandas as pddf = pd.DataFrame(history)plt.figure(figsize=(12, 8))plt.plot(	df['epoch'],	df['train_loss'],	label='Training Loss',	marker='o',	linestyle='-')plt.plot(	df['epoch'],	df['test_loss'],	label='Validation Loss',	marker='o',	linestyle='--')plt.title('Train/Validation Loss over Epochs')plt.xlabel('Epoch')plt.ylabel('Loss')plt.legend()plt.grid(True)plt.show()

这将返回一个显示每个时期损失的图表。损失应随着时间的推移而减少。

我们的训练笔记本还通过一个示例展示了如何计算各个时期的 AP。

测试你的 RF-DETR 模型

训练好的模型准备好后，下一步就是对示例图像进行推理。

为了从我们的测试集中加载数据并可视化示例图像上的预测，我们将使用 Supervisor Python 包。Supervisor 有一系列可用于构建计算机视觉应用程序的实用程序，包括数据加载器和注释器。

让我们加载测试集：

import supervision as svds = sv.DetectionDataset.from_coco(	images_directory_path=f"{dataset.location}/test",	annotations_path=f"{dataset.location}/test/_annotations.coco.json",)

接下来，让我们从测试集中加载一个随机图像，并将真实情况与模型的结果进行比较：

path, image, annotations = ds[4]from rfdetr import RFDETRBasefrom rfdetr.util.coco_classes import COCO_CLASSESimport supervision as svimport numpy as npfrom PIL import Imageimage = Image.open(path)detections = model.predict(image, threshold=0.5)text_scale = sv.calculate_optimal_text_scale(resolution_wh=image.size)thickness = sv.calculate_optimal_line_thickness(resolution_wh=image.size)bbox_annotator = sv.BoxAnnotator(thickness=thickness)label_annotator = sv.LabelAnnotator(	text_color=sv.Color.BLACK,	text_scale=text_scale,	text_thickness=thickness,	smart_position=True)annotations_labels = [	f"{ds.classes[class_id]}"	for class_id	in annotations.class_id]detections_labels = [	f"{ds.classes[class_id]} {confidence:.2f}"	for class_id, confidence	in zip(detections.class_id, detections.confidence)]annotation_image = image.copy()annotation_image = bbox_annotator.annotate(annotation_image, annotations)annotation_image = label_annotator.annotate(annotation_image, annotations, annotations_labels)detections_image = image.copy()detections_image = bbox_annotator.annotate(detections_image, detections)detections_image = label_annotator.annotate(detections_image, detections, detections_labels)sv.plot_images_grid(images=[annotation_image, detections_image], grid_size=(1, 2), titles=["Annotation", "Detection"])