在自定义数据上训练 YOLOv5

📚 本指南介绍了如何使用 YOLOv5 模型训练你自己的 自定义数据集 🚀。训练自定义模型是将 计算机视觉 解决方案应用于通用 目标检测 之外的特定实际应用的关键步骤。

在你开始之前

首先,请确保你已经设置好了必要的环境。克隆 YOLOv5 仓库并安装 requirements.txt 中要求的依赖项。一个 Python>=3.8.0 环境和 PyTorch>=1.8 是必不可少的。如果本地未找到模型和数据集,它们将自动从最新的 YOLOv5 发布版本 中下载。

git clone https://github.com/ultralytics/yolov5 # Clone the repository
cd yolov5
pip install -r requirements.txt # Install dependencies

在自定义数据上训练

Ultralytics 主动学习循环图

开发自定义 目标检测 模型是一个迭代过程:

  1. 收集并整理图像:收集与你特定任务相关的图像。高质量、多样化的数据至关重要。请参阅我们的 数据收集与标注 指南。
  2. 标注对象:准确标注图像中你感兴趣的对象。
  3. 训练模型:使用标注数据 训练 你的 YOLOv5 模型。通过使用预训练权重来利用 迁移学习
  4. 部署与预测:利用训练好的模型对新的、未见过的数据进行 推理
  5. 收集边缘情况:识别模型表现不佳的场景(边缘情况),并将类似数据添加到你的数据集中以提高稳健性。重复此循环。

Ultralytics Platform 为整个 机器学习运维 (MLOps) 周期提供了一个精简的免代码解决方案,包括数据集管理、模型训练和部署。

许可协议

Ultralytics 提供两种许可选项以适应不同的使用场景:

  • AGPL-3.0 许可证:这种 OSI 认证 的开源许可证非常适合热衷于开放协作和知识共享的学生、研究人员和爱好者。它要求衍生作品必须在同一许可证下共享。查看 LICENSE 文件以了解全部详情。
  • 企业许可证:专为商业应用而设计,该许可证允许将 Ultralytics 软件和 AI 模型无缝集成到商业产品和服务中,而无需遵守 AGPL-3.0 的开源规定。如果你的项目需要商业部署,请申请 企业许可证

Ultralytics 许可 页面进一步探索我们的许可选项。

在开始训练之前,数据集准备工作必不可少。

创建数据集

YOLOv5 模型需要标注数据来学习对象类别的视觉特征。正确组织数据集是关键。

1.1 创建 dataset.yaml

数据集配置文件(例如 coco128.yaml)概述了数据集的结构、类名和图像目录路径。COCO128 是一个小型的示例数据集,包含来自庞大 COCO 数据集的前 128 张图像。它对于快速测试训练流程和诊断潜在问题(如 过拟合)非常有用。

dataset.yaml 文件结构包括:

  • path:包含数据集的根目录。
  • train, val, test:从 path 到包含图像的目录的相对路径,或者列出用于训练、验证和测试集的图像路径的文本文件。
  • names:将类索引(从 0 开始)映射到其对应类名的字典。

当你从 YOLOv5 仓库根目录启动训练时,你可以将 path 设置为绝对目录(例如 /home/user/datasets/coco128)或相对路径(如 ../datasets/coco128)。

以下是 coco128.yaml 的结构(在 GitHub 上查看):

# Dataset root directory relative to the yolov5 directory
path: coco128

# Train/val/test sets: specify directories, *.txt files, or lists
train: images/train2017 # 128 images for training
val: images/train2017 # 128 images for validation
test: # Optional path to test images

# Classes (example using 80 COCO classes)
names:
    0: person
    1: bicycle
    2: car
    # ... (remaining COCO classes)
    77: teddy bear
    78: hair drier
    79: toothbrush

1.2 利用模型进行自动化标注

手动标注是常见的方法,但非常耗时。基础模型可以实现标注的自动化或半自动化,从而加快数据集的创建速度。可以帮助生成标签的模型示例:

  • Google Gemini:像 Gemini 这样的大型多模态模型拥有强大的图像理解能力。可以通过提示它们来识别和定位图像中的对象,生成可以转换为 YOLO 格式标签的边界框或描述。在提供的教程笔记本中探索其潜力。
  • SAM2 (Segment Anything Model 2):像 SAM2 这样专注于分割的基础模型,可以高精度地识别和描绘对象。虽然主要用于分割,但得到的掩码通常可以转换为适用于目标检测任务的边界框标注。
  • YOLOWorld:该模型提供了开放词汇检测能力。你可以提供你感兴趣对象的文本描述,YOLOWorld 可以在图像中定位它们,而无需针对这些特定类别进行预先训练。这可以用作生成初始标签的起点,然后可以进行优化。

使用这些模型可以提供“预标注”步骤,减少所需的手动工作。但是,审查和优化自动生成的标签以确保准确性和一致性至关重要,因为质量直接影响你训练的 YOLOv5 模型的性能。生成(并可能优化)标签后,请确保它们遵守 YOLO 格式:每张图像一个 *.txt 文件,每行代表一个对象为 class_index x_center y_center width height(归一化坐标,从 0 开始计数的类别)。如果图像中没有感兴趣的对象,则不需要对应的 *.txt 文件。

YOLO 格式 *.txt 文件的规范很精确:

  • 每个 边界框 一行。
  • 每行必须包含:class_index x_center y_center width height
  • 坐标必须 归一化 到 0 到 1 之间的范围。为此,请将 x_centerwidth 的像素值除以图像的总宽度,并将 y_centerheight 除以图像的总高度。
  • 类索引是从 0 开始计数的(即第一个类别由 0 表示,第二个由 1 表示,依此类推)。

Example image with two persons and a tie annotated

上图对应的标签文件,包含两个 'person' 对象(类索引 0)和一个 'tie' 对象(类索引 27),看起来像这样:

YOLO format label file content example

1.3 组织目录

Structure your datasets directory as illustrated below. By default, YOLOv5 anticipates the dataset directory (e.g., /coco128) to reside within a /datasets folder located adjacent to the /yolov5 repository directory.

YOLOv5 会自动通过将图像路径中最后一个 /images/ 实例替换为 /labels/ 来定位每张图像的标签。例如:

../datasets/coco128/images/im0.jpg # Path to the image file
../datasets/coco128/labels/im0.txt # Path to the corresponding label file

推荐的目录结构是:

/datasets/
└── coco128/  # Dataset root
    ├── images/
    │   ├── train2017/  # Training images
    │   │   ├── 000000000009.jpg
    │   │   └── ...
    │   └── val2017/    # Validation images (optional if using same set for train/val)
    │       └── ...
    └── labels/
        ├── train2017/  # Training labels
        │   ├── 000000000009.txt
        │   └── ...
        └── val2017/    # Validation labels (optional if using same set for train/val)
            └── ...

YOLOv5 recommended dataset directory structure

选择模型

选择一个 预训练模型 来开始训练过程。与从头开始训练相比,从预训练权重开始可以显著加快学习速度并提高性能。YOLOv5 提供各种模型大小,每种都在速度和准确性之间进行不同的权衡。例如,YOLOv5s 是第二小且速度最快的模型,适用于资源受限的环境。请参阅 README 表格 以获取所有可用 模型 的详细比较。

Comparison chart of YOLOv5 models showing size, speed, and accuracy

训练

Begin the model training using the train.py script. Essential arguments include:

  • --img:定义输入 图像尺寸(例如 --img 640)。较大的尺寸通常会产生更好的准确性,但需要更多的 GPU 内存。
  • --batch:确定 批量大小(例如 --batch 16)。选择你的 GPU 能处理的最大尺寸。
  • --epochs:指定训练 轮次 的总数(例如 --epochs 100)。一个轮次代表对整个训练数据集进行一次完整遍历。
  • --data:指向你 dataset.yaml 文件的路径(例如 --data coco128.yaml)。
  • --weights:指向初始权重文件的路径。强烈建议使用预训练权重(例如 --weights yolov5s.pt),以实现更快的收敛和更好的结果。若要从头开始训练(除非你有非常大的数据集和特定需求,否则不建议),请使用 --weights '' --cfg yolov5s.yaml

如果本地未找到预训练权重,它们将自动从 最新的 YOLOv5 发布版本 中下载。

# Example: Train YOLOv5s on the COCO128 dataset for 3 epochs
python train.py --img 640 --batch 16 --epochs 3 --data coco128.yaml --weights yolov5s.pt
优化训练速度

💡 使用 --cache ram--cache disk 将数据集图像分别缓存到 RAM 或本地磁盘中。这会显著加速训练,特别是在数据集 I/O(输入/输出)操作成为瓶颈时。请注意,这需要大量的 RAM 或磁盘空间。

本地数据存储

💡 始终使用存储在本地的数据集进行训练。从网络驱动器(如 Google Drive)或远程存储访问数据可能会明显变慢,并阻碍训练性能。将数据集复制到本地 SSD 通常是最佳实践。

所有训练输出(包括权重和日志)都会保存到 runs/train/ 目录中。每次训练会话都会创建一个新的子目录(例如 runs/train/exp, runs/train/exp-2 等)。如需交互式、亲手实践的体验,请查看我们官方教程笔记本中的训练部分:Open In Colab Open In Kaggle

可视化

YOLOv5 与各种用于可视化训练进度、评估结果和实时监控性能的工具无缝集成。

Comet 日志记录与可视化

Comet 已完全集成,可进行全面的实验跟踪。实时可视化指标、保存超参数、管理数据集和模型检查点,并使用交互式 Comet 自定义面板 分析模型预测。

入门非常简单:

pip install comet_ml                                                          # 1. Install Comet library
export COMET_API_KEY=YOUR_API_KEY_HERE                                        # 2. Set your Comet API key (create a free account at Comet.ml)
python train.py --img 640 --epochs 3 --data coco128.yaml --weights yolov5s.pt # 3. Train your model - Comet automatically logs everything!

在我们的 Comet 集成指南 中深入了解支持的功能。从他们的官方 文档 中了解更多关于 Comet 的功能。尝试 Comet Colab 笔记本获取现场演示:Open In Colab

Comet UI showing YOLOv5 training metrics and visualizations

ClearML 日志记录与自动化

ClearML 集成支持详细的实验跟踪、数据集版本管理,甚至训练运行的远程执行。通过以下简单的步骤激活 ClearML:

  • 安装包:pip install clearml
  • 初始化 ClearML:运行 clearml-init 一次以连接到你的 ClearML 服务器(无论是自托管还是 免费层级)。

ClearML 会自动捕获实验详情、模型上传、比较、未提交的代码更改以及已安装的包,确保完全的可重复性。你可以使用 ClearML Data 轻松安排远程代理上的训练任务并管理数据集版本。探索 ClearML 集成指南 以获取全面详细信息。

ClearML experiment management UI for YOLOv5

本地日志记录

训练结果会使用 TensorBoard 自动记录,并以 CSV 文件的形式保存在特定的实验目录中(例如 runs/train/exp)。记录的数据包括:

  • 训练和验证损失以及性能指标。
  • 显示所应用增强(如马赛克增强)的示例图像。
  • 地面实况标签与模型预测并排,以供目测。
  • 关键评估指标,如 精确率-召回率 (PR) 曲线。
  • 混淆矩阵,用于详细的类级别性能分析。
YOLOv5 local logging results with charts and mosaics

results.csv 文件会在每一轮后更新,并在训练结束后绘制为 results.png。你也可以使用提供的实用函数手动绘制任何 results.csv 文件:

from utils.plots import plot_results

# Plot results from a specific training run directory
plot_results("runs/train/exp/results.csv")  # This will generate 'results.png' in the same directory

YOLOv5 results.png training metrics plot

后续步骤

训练成功完成后,表现最佳的模型检查点 (best.pt) 会被保存,并准备好进行部署或进一步优化。可能的后续步骤包括:

支持的环境

Ultralytics 提供预装了必要依赖项(如 CUDA, cuDNN, Python, 和 PyTorch)的即用型环境,有助于顺利起步。

项目状态

YOLOv5 持续集成状态徽章

此徽章表示所有 YOLOv5 GitHub Actions 持续集成 (CI) 测试均已成功通过。这些严格的 CI 测试涵盖了核心功能,包括在 macOS、Windows 和 Ubuntu 操作系统上的 训练验证推理导出基准测试。测试每 24 小时以及每次代码提交时都会自动执行,从而确保一致的稳定性和最佳性能。

常见问题 (FAQ)

我该如何在我的自定义数据集上训练 YOLOv5?

在自定义数据集上训练 YOLOv5 涉及几个关键步骤:

  1. 准备你的数据集:收集图像并进行标注。确保标注采用所需的 YOLO format。将图像和标签组织到 train/val/(以及可选的 test/)目录中。考虑使用 Google GeminiSAM2YOLOWorld 等模型来辅助或自动化标注过程(请参阅第 1.2 节)。
  2. 设置你的环境:克隆 YOLOv5 仓库,并使用 pip install -r requirements.txt 安装依赖项。
    git clone https://github.com/ultralytics/yolov5
    cd yolov5
    pip install -r requirements.txt
  3. 创建数据集配置:在 dataset.yaml 文件中定义数据集路径、类别数量和类名。
  4. 开始训练:执行 train.py 脚本,提供你的 dataset.yaml 路径、所需的预训练权重(例如 yolov5s.pt)、图像大小、批处理大小以及 epoch 数量。
    python train.py --img 640 --batch 16 --epochs 100 --data path/to/your/dataset.yaml --weights yolov5s.pt

为什么我应该使用 Ultralytics Platform 来训练我的 YOLO 模型?

Ultralytics Platform 是一个综合性平台,旨在简化整个 YOLO 模型开发生命周期,通常无需编写任何代码。主要优势包括:

  • 简化训练:使用预配置的环境和直观的用户界面轻松训练模型。
  • 集成数据管理:在平台内高效上传、版本控制和管理你的数据集。
  • 实时监控:使用 Comet 或 TensorBoard 等集成工具跟踪训练进度并可视化性能指标。
  • 协作功能:通过共享资源、项目管理工具和轻松的模型共享来促进团队合作。
  • 无代码部署:将训练好的模型直接部署到各种目标设备。

如需实操演示,请查看我们的博客文章:如何使用 Ultralytics Platform 训练你的自定义模型

我该如何将已标注的数据转换为 YOLOv5 格式?

无论你是手动标注还是使用自动工具(如第 1.2 节中提到的那些),最终的标签必须采用 YOLOv5 所需的特定 YOLO format

  • Create one .txt file for each image. The filename should match the image filename (e.g., image1.jpg corresponds to image1.txt). Place these files in a labels/ directory parallel to your images/ directory (e.g., ../datasets/mydataset/labels/train/).
  • .txt 文件中的每一行代表一个对象标注,并遵循以下格式:class_index center_x center_y width height
  • 坐标(center_xcenter_ywidthheight)必须相对于图像尺寸进行归一化(值在 0.0 到 1.0 之间)。
  • 类别索引是从零开始的(第一个类别为 0,第二个为 1,依此类推)。

许多手动标注工具提供直接导出为 YOLO 格式的功能。如果使用自动模型,你需要编写脚本或流程将它们的输出(例如边界框坐标、分割掩码)转换为这种特定的归一化文本格式。确保你的最终数据集结构符合指南中提供的示例。欲了解更多详情,请参阅我们的 数据收集和标注指南

在商业应用中使用 YOLOv5 有哪些许可选项?

Ultralytics 提供针对不同需求量身定制的灵活许可方案:

  • AGPL-3.0 License:此开源许可适用于学术研究、个人项目以及可接受开源合规的情况。它要求修改和衍生作品也必须在 AGPL-3.0 下开源。请查看 AGPL-3.0 License 详情
  • Enterprise License: A commercial license designed for businesses integrating YOLOv5 into proprietary products or services. This license removes the open-source obligations of AGPL-3.0, allowing for closed-source distribution. Visit our Licensing page for further details or to request an Enterprise License.

选择最符合你项目要求和分发模型的许可。

评论