在自定义数据上训练 YOLOv5#

📚 本指南介绍了如何使用 YOLOv5 模型训练你自己的自定义数据集 🚀。训练自定义模型是根据通用目标检测之外的具体现实应用来定制计算机视觉解决方案的基础步骤。

开始之前#

首先，确保配好所需的环境。克隆 YOLOv5 仓库并从 requirements.txt 安装所需的依赖项。具备 PyTorch>=1.8 的 Python>=3.8.0 环境至关重要。如果在本地找不到模型和数据集，系统将自动从最新的 YOLOv5 发布版本中下载它们。

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

在自定义数据上进行训练#

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

1. 收集和整理图像：搜集与你的特定任务相关的图像。高质量、多样化的数据至关重要。请参阅我们的数据收集与标注指南。
2. 标注对象：准确标注你图像中感兴趣的对象。
3. 训练模型：使用标注好的数据来训练你的 YOLOv5 模型。通过从预训练权重开始，充分利用迁移学习。
4. 部署与预测：利用训练好的模型对新的、未见过的数据进行推理。
5. 收集边缘案例：找出模型表现不佳的场景（边缘案例），并将类似的数据添加到你的数据集中以提高鲁棒性。重复此循环。

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

开始训练之前，数据集的准备工作至关重要。

创建数据集#

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（归一化坐标，以零为索引的类别）。如果图片中没有感兴趣的对象，则不需要相应的 *.txt 文件。

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

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

对应于上图的标签文件包含两个“person”对象（类别索引 0）和一个“tie”对象（类别索引 27），其外观如下：

1.3 组织目录#

按照下图所示组织你的数据集目录。默认情况下，YOLOv5 预期数据集目录（例如 /coco128）位于与 /yolov5 仓库目录相邻的 /datasets 文件夹中。

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 提供了各种模型大小，每种模型都在速度和准确率之间有着不同的平衡。例如，YOLOv5s 是第二小且速度最快的模型，适合资源受限的环境。请查阅 README 表格了解所有可用模型的详细对比。

训练#

使用 train.py 脚本开始模型训练。必要的参数包括：

• --img：定义输入图像大小（例如 --img 640）。更大的尺寸通常会产生更好的准确性，但需要更多的 GPU 内存。
• --batch：确定批次大小（例如 --batch 16）。选择你的 GPU 可以处理的最大尺寸。
• --epochs：指定训练轮数的总数（例如 --epochs 100）。一个 epoch 表示对整个训练数据集的完整遍历。
• --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

所有训练输出（包括权重和日志）都保存在 runs/train/ 目录中。每个训练会话都会创建一个新的子目录（例如 runs/train/exp、runs/train/exp2 等）。要获得交互式的动手体验，请浏览我们官方教程笔记本中的训练部分：

可视化#

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 笔记本以获取实时演示：

ClearML 日志记录与自动化#

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

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

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

本地日志记录#

训练结果使用 TensorBoard 自动记录，并作为 CSV 文件保存在特定实验目录中（例如 runs/train/exp）。记录的数据包括：

• 训练和验证损失以及性能指标。
• 显示应用增强（如马赛克增强）的示例图像。
• 用于视觉检查的真实标签与模型预测。
• 关键评估指标，例如精确率-召回率 (PR) 曲线。
• 用于详细的逐类性能分析的混淆矩阵。

results.csv 文件在每个 epoch 后更新，并在训练结束后绘制为 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

后续步骤#

训练成功完成后，表现最佳的模型检查点 (best.pt) 将被保存，随时可用于部署或进一步完善。潜在的下一步包括：

• 通过 CLI 或 Python 使用训练好的模型对新图像或视频运行推理。
• 执行验证以评估模型在不同数据分割（例如保留的测试集）上的准确率和泛化能力。
• 将模型导出为各种部署格式，例如 ONNX、TensorFlow SavedModel 或 TensorRT，以便在不同平台上进行优化的推理。
• 采用超参数调优技术以潜在地榨取额外的性能提升。
• 遵循我们的最佳训练结果提示，并根据性能分析迭代添加更多样化和具有挑战性的数据，继续改进你的模型。

支持的环境#

Ultralytics 提供配备了 CUDA、cuDNN、Python 和 PyTorch 等基本依赖项的即用型环境，从而促进顺利起步。

• 免费 GPU 笔记本：
  • 
  • 
  • 
• 云平台：
  • Google Cloud：GCP 快速入门指南
  • Amazon AWS：AWS 快速入门指南
  • Microsoft Azure：AzureML 快速入门指南
• 本地设置：
  • Docker：Docker 快速入门指南

项目状态#

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

常见问题解答#