使用自定义数据训练 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

使用自定义数据训练#

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

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

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

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

1. 创建数据集#

YOLOv5 模型需要带标签的数据来学习对象类别的视觉特征。正确整理数据集是关键。

1.1 创建 dataset.yaml#

数据集配置文件（例如 coco128.yaml）描述数据集的结构、类别名称以及图像目录的路径。COCO128 是一个小型示例数据集，由大型 COCO 数据集中的前 128 张图像组成。它适合用于快速测试训练流程，以及诊断过拟合等潜在问题。

dataset.yaml 文件结构包括：

• path：包含数据集的根目录。
• train、val、test：从 path 到包含图像的目录，或包含训练集、验证集和测试集图像路径的文本文件的相对路径。
• names：将类别索引（从 0 开始）映射到相应类别名称的字典。

启动训练时，你可以将 path 设置为绝对目录（例如 /home/user/datasets/coco128），也可以设置为相对路径（例如 ../datasets/coco128），前提是从 YOLOv5 代码仓库根目录启动。

下面是 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（任意分割模型 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 表示，第二个类别由 1 表示，以此类推）。

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

1.3 整理目录#

按照下方所示整理你的数据集目录。默认情况下，YOLOv5 预期数据集目录（例如 /coco128）位于 /datasets 文件夹内，而该文件夹与 /yolov5 代码仓库目录处于同级。

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)
            └── ...

2. 选择模型#

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

3. 训练#

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

• --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

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

4. 可视化#

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 文件会在每个周期结束后更新，并在训练完成后绘制为 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

5. 后续步骤#

训练成功完成后，性能最佳的模型检查点（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 小时自动执行一次，并在每次代码提交时执行，从而确保稳定性和最佳性能的一致性。

常见问题#