YOLOv12训练自己的数据集——利用Google Colab环境和Roboflow工具

YOLOv12训练自己的数据集——利用Google Colab环境和Roboflow工具

引言

2025年2月18日，YOLOv12作为一种前沿的计算机视觉模型架构正式发布，由布法罗大学和中科院大学的Yunjie Tian、Qixiang Ye和David Doermann在论文《YOLOv12: Attention-Centric Real-Time Object Detectors》中提出。YOLOv12以注意力机制为核心设计，在Microsoft COCO数据集上实现了更低的推理延迟和更高的均值平均精度（mAP），成为实时目标检测领域的又一里程碑。官方代码已开源于 GitHub，支持用户对其进行微调和训练。

论文详解请参考这篇文章《YOLOv12 - 注意力机制驱动的实时目标检测新篇章》——论文详解

YOLOv12源码及权重下载——YOLOv12

本文将详细介绍如何在自定义数据集上训练YOLOv12模型，包括数据集准备、模型训练和推理验证的全流程。您将学习如何利用Google Colab环境和Roboflow工具，基于一个示例数据集（集装箱检测）完成YOLOv12的训练，并获得可用于实际应用的目标检测模型。

YOLOv12简介

YOLOv12通过引入区域注意力（Area Attention）和残差高效层聚合网络（R-ELAN），显著提高了检测精度，同时保持了实时性。以下是其在COCO数据集上的主要性能指标（见论文Table 1）：

这些结果表明，YOLOv12在不同规模下均优于YOLOv11、YOLOv10等前代模型，且相较于RT-DETR等端到端检测器，计算量和参数量更少，速度更快。

训练准备

在开始训练之前，您需要准备自定义数据集并配置训练环境。以下是具体步骤：

步骤1：安装YOLOv12及依赖项

我们将在Google Colab中运行训练过程。打开Roboflow提供的YOLOv12训练笔记本（可从GitHub或Roboflow获取完整代码），并按以下步骤安装依赖：

# 克隆YOLOv12仓库
!git clone https://github.com/sunsmarterjie/yolov12
%cd yolov12

# 安装必要的Python包
!pip install roboflow supervision flash-attn --upgrade -q
!pip install -r requirements.txt
!pip install -e .
!pip install --upgrade flash-attn

• YOLOv12：核心模型代码，目前需从源代码构建，因尚未发布pip包。
• supervision：用于后处理模型预测结果。
• Roboflow：用于下载和管理数据集。
• flash-attn：优化注意力机制的计算效率，YOLOv12依赖此库。

步骤2：准备自定义数据集

您需要一个带有标注的目标检测数据集。如果没有现成数据集，可参考Roboflow的《Getting Started with Roboflow》指南创建。本文以开源的集装箱检测数据集（Yard Management System, version 11）为例：

1. 下载数据集：
   • 创建免费Roboflow账户并获取API密钥。
   • 在Colab中运行以下代码下载数据集：

from google.colab import userdata
from roboflow import Roboflow

ROBOFLOW_API_KEY = userdata.get('ROBOFLOW_API_KEY')
rf = Roboflow(api_key=ROBOFLOW_API_KEY)
project = rf.workspace("roboflow-universe-projects").project("yard-management-system")
version = project.version(11)
dataset = version.download("yolov8")

2. 调整数据集格式：
   • YOLOv12使用与YOLOv8兼容的TXT标注格式，但需要修改data.yaml文件以适配路径。运行以下代码：

!sed -i '$d' {dataset.location}/data.yaml
!sed -i '$d' {dataset.location}/data.yaml
!sed -i '$d' {dataset.location}/data.yaml
!sed -i '$d' {dataset.location}/data.yaml
!echo -e "test: ../test/images\ntrain: ../train/images\nval: ../valid/images" >> {dataset.location}/data.yaml

此步骤移除原始文件中可能不兼容的行，并添加正确的训练、验证和测试集路径。

模型训练

步骤3：训练YOLOv12模型

配置好数据集后，我们可以开始训练。YOLOv12提供五种预训练权重（yolov12n.pt、yolov12s.pt、yolov12m.pt、yolov12l.pt、yolov12x.pt），可根据计算资源选择。本例使用yolov12s.pt，建议至少训练250个epoch。

运行以下代码：

from ultralytics import YOLO

model = YOLO('yolov12s.yaml')
results = model.train(
    data=f'{dataset.location}/data.yaml',
    epochs=250
)

• 参数说明：
  • data：指向调整后的data.yaml文件。
  • epochs：训练轮数，论文建议600个epoch，但250个epoch已可获得较好结果。
  • 默认使用SGD优化器，初始学习率0.01，线性衰减至0.0001（见论文Table 7）。

注意：训练完成后，model()可能抛出TypeError: argument of type 'PosixPath' is not iterable错误（已提交GitHub issue），但模型权重仍会保存，可继续后续步骤。

训练结果评估

训练完成后，可查看混淆矩阵和训练曲线：

1. 混淆矩阵：

from IPython.display import Image
Image(filename=f'/content/yolov12/runs/detect/train/confusion_matrix.png', width=600)

2. 训练曲线：

Image(filename=f'/content/yolov12/runs/detect/train/results.png', width=600)

推理验证

步骤4：使用训练好的模型进行推理

加载训练后的最佳权重，在验证集上运行推理：

import random
import cv2
import supervision as sv

model = YOLO(f"/content/yolov12/runs/detect/train/weights/best.pt")

# 加载验证集
ds = sv.DetectionDataset.from_yolo(
    images_directory_path=f"{dataset.location}/valid/images",
    annotations_directory_path=f"{dataset.location}/valid/labels",
    data_yaml_path=f"{dataset.location}/data.yaml"
)

# 随机选择一张图片进行推理
image_path = random.choice(list(ds.images.keys()))
image = cv2.imread(image_path)

# 运行模型并后处理
results = model(image)[0]
detections = sv.Detections.from_ultralytics(results).with_nms()

# 添加边界框和标签
bounding_box_annotator = sv.BoundingBoxAnnotator()
label_annotator = sv.LabelAnnotator()
annotated_image = bounding_box_annotator.annotate(scene=image, detections=detections)
annotated_image = label_annotator.annotate(scene=annotated_image, detections=detections)

# 显示结果
sv.plot_image(annotated_image)

结果示例：
在集装箱数据集上，模型成功识别出集装箱、ID、标志和底盘ID四个类别。

结论

YOLOv12凭借其注意力中心的架构设计，为实时目标检测提供了新的可能性。本文通过一个集装箱检测示例，展示了如何在自定义数据集上训练YOLOv12，包括环境配置、数据准备、模型训练和推理验证。完整代码已作为笔记本开源，可从Roboflow或GitHub获取。

未来，您可以尝试调整超参数（如论文Table 7中的数据增强配置），或使用更大规模的模型（如YOLOv12-X）以进一步提升性能。更多训练教程可参考Roboflow Notebooks。