跳转至内容

在自定义数据上训练 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 HUB为此整个机器学习运维 (MLOps)周期提供了一个简化的、无需代码的解决方案,包括数据集管理、模型训练和部署。

许可

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

  • AGPL-3.0 许可证: 这种经 OSI 批准的开源许可证非常适合热衷于开放协作和知识共享的学生、研究人员和爱好者。它要求衍生作品以相同的许可证共享。有关完整详细信息,请参见LICENSE文件。
  • 企业许可证: 此许可证专为商业应用而设计,允许将 Ultralytics 软件和 AI 模型无缝集成到商业产品和服务中,而无需 AGPL-3.0 的开源规定。 如果您的项目需要商业部署,请申请企业许可证

Ultralytics 许可 页面上进一步了解我们的许可选项。

在开始训练之前,数据集准备至关重要。

1. 创建数据集

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

1.1 创建 dataset.yaml

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

字段 dataset.yaml 文件结构包括:

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

以下是结构: 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 格式标签的边界框或描述。请在提供的教程 notebook 中探索其潜力。
  • SAM2 (第二代分割一切模型): 专注于分割的基础模型,如 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_centerwidth 除以图像的总宽度,然后除以 y_centerheight 除以图像的总高度。
  • 类索引从零开始(即,第一个类由以下表示 0,第二个通过 1,等等)。

带有两个人和一条领带的注释示例图像

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

带标注图像的 YOLO 格式标签文件内容示例

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

图示显示了推荐的 YOLOv5 数据集目录结构

2. 选择模型

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

YOLOv5 模型的大小、速度和准确性比较图表

3. 训练

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

  • --img: 定义输入 图像大小 (例如, --img 640)。较大的尺寸通常会产生更好的精度,但需要更多的 GPU 内存。
  • --batch: 确定 批次大小 (例如, --batch 16)。选择您的 GPU 可以处理的最大尺寸。
  • --epochs: 指定训练的总轮数 个 epoch (例如, --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

优化训练速度

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

本地数据存储

💡 始终使用本地存储的数据集进行训练。从网络驱动器(如 Google Drive)或远程存储访问数据可能会显著降低速度并影响训练性能。将数据集复制到本地 SSD 通常是最佳实践。

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

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 功能的信息 文档。请试用 Comet Colab Notebook 以获取实时演示: 在 Colab 中打开

Comet UI 显示 YOLOv5 训练指标和可视化效果

ClearML 日志记录和自动化 🌟 新功能

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

  • 安装软件包: pip install clearml
  • 初始化 ClearML:运行 clearml-init 只需一次即可连接到您的 ClearML 服务器(自托管或 免费层)。

ClearML 自动捕获实验详细信息、模型上传、比较、未提交的代码更改和已安装的软件包,确保完全可重现性。您可以使用 ClearML Data 轻松地在远程代理上安排训练任务并管理数据集版本。请浏览 ClearML 集成指南 以获取全面的详细信息。

ClearML 实验管理用户界面,显示 YOLOv5 训练运行的图表和日志

本地日志记录

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

  • 训练和验证损失以及性能指标。
  • 显示应用增强(如马赛克)的 sample_images。
  • 与模型预测一起显示的 Ground Truth 标签,用于进行可视化检查。
  • 关键评估指标,例如精确率-召回率 (PR) 曲线。
  • 混淆矩阵,用于详细的类性能分析。

来自 YOLOv5 训练的本地日志记录结果示例,包括图表和图像马赛克

字段 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

示例 results.png 图,显示了训练指标,如 mAP、精确率、召回率以及随 epoch 变化的损失

5. 下一步

成功完成训练后,性能最佳的模型检查点(best.pt)已保存,可以部署或进一步优化。潜在的后续步骤包括:

  • 通过 CLIPython,使用训练好的模型对新图像或视频运行推理
  • 执行验证,以评估模型在不同数据分割(例如,预留的测试集)上的准确性和泛化能力。
  • 将模型导出为各种部署格式,如 ONNXTensorFlow SavedModelTensorRT,以便在不同的平台上进行优化推理。
  • 采用超参数调整技术来潜在地挤出额外的性能提升。
  • 参考我们的最佳训练结果技巧,并根据性能分析,迭代地添加更多样化和具有挑战性的数据,以持续改进您的模型。

支持的环境

Ultralytics 提供了即用型环境,配备了必要的依赖项,如 CUDAcuDNNPythonPyTorch,从而方便您顺利启动。

项目状态

YOLOv5 持续集成状态徽章

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

常见问题

常见问题

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

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

  1. 准备您的数据集: 收集图像并进行标注。确保标注符合要求的格式。 YOLO 格式。将图像和标签整理到 train/val/ (并且可以选择 test/)目录。考虑使用像 Google Gemini, SAM2YOLOWorld 以协助或自动化标注过程(参见第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 HUB 来训练我的 YOLO 模型?

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

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

有关实际演练,请查看我们的博客文章:如何使用 Ultralytics HUB 训练您的自定义模型

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

无论您是手动标注还是使用自动化工具(如第 1.2 节中提到的工具),最终标签都必须采用 YOLOv5 要求的特定 YOLO 格式

  • 创建一个 .txt 每个图像的文件。文件名应与图像文件名匹配(例如, image1.jpg 对应于 image1.txt)。将这些文件放在 labels/ 目录,与您的 images/ 目录(例如, ../datasets/mydataset/labels/train/)。
  • 在每一行中: .txt 文件代表一个对象注释,格式如下: class_index center_x center_y width height.
  • 坐标 (center_x, center_y, width, height)必须是 已标准化 (介于 0.0 和 1.0 之间的值)相对于图像的尺寸。
  • 类索引是 基于零 (第一个类别是 0,第二个是 1,等等)。

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

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

Ultralytics 提供了灵活的许可,以满足不同的需求:

  • AGPL-3.0 许可证: 此开源许可证适用于学术研究、个人项目以及可以接受开源合规性的情况。它强制要求修改和衍生作品也必须在 AGPL-3.0 下开源。请查看AGPL-3.0 许可证详细信息
  • 企业许可证: 一种商业许可证,专为将 YOLOv5 集成到专有产品或服务中的企业而设计。 此许可证消除了 AGPL-3.0 的开源义务,允许闭源分发。 请访问我们的许可页面了解更多详情或申请企业许可证

选择最符合您的项目需求和分发模式的许可证。



📅 创建于 1 年前 ✏️ 更新于 2 个月前

评论