如何训练自定义 YOLOv5 模型来检测对象

• 下载“critters”模型 - 12.5 MB
• 下载自定义模型训练 Jupyter Notebook - 110.4 KB

引言

目标检测是人工智能的一个常见应用。在 CodeProject.AI Server 中，我们添加了一个使用 YOLOv5 架构进行目标检测的模块。这是一个出色的系统，但使用“标准”YOLOv5 模型意味着您仅限于默认模型中提供的 80 个类别。要检测此默认集之外的对象，您需要训练自己的自定义模型。

虽然 YOLOv5 文档会引导您完成创建新模型的过程，但设置训练会话以及生成良好且准确的模型并不像他们认为的那么简单。

本文将通过以下步骤引导您创建自定义模型以检测后院害虫

1. 设置您的训练环境。
2. 获取足够大且经过标注的相关且多样化的图像集，用于模型的训练、测试和验证。
3. 训练您的模型，包括启用 GPU 加速的关键步骤
4. 在 CodeProject.AI Server 中使用您的自定义模型

我观察到的一些训练过程出乎意料，我计划在探索和改进过程以及生成的模型时扩展本文。

有许多术语和指标用于描述 AI 模型的性能。我将使用这些指标来确定训练好的模型对于其训练任务来说有多“好”。您可以在 关于训练指标的说明 中阅读有关这些内容的说明。

设置您的训练环境

硬件

我用于 AI 开发和测试的系统具有以下规格

开发 IDE

我将 Visual Studio Code 作为我的开发 IDE，因为它可以在 Windows 和 Linux 上运行。我已经 配置好了 Python 开发，并使用 Python Jupyter Notebook 来执行和记录结果。我正在运行 Python 3.9.13。

设置项目文件夹

对于我的项目，我创建了一个目录 c:\Dev\YoloV5_Training 并在 Visual Studio Code 中打开了它。然后，我创建了一个名为 Custom Model Training.ipynb 的 Python Notebook。完成这些之后，让我们开始吧。

设置 Python 虚拟环境

为了不污染全局 Python 安装，我将使用 Python venv 命令为该项目创建一个虚拟环境。此虚拟环境将创建在 venv 子目录中。

仅运行一次以下命令。

!Python -m venv venv

完成此操作后，您需要告诉 Visual Studio Code 您想使用您刚刚创建的虚拟环境。您可以通过单击 Notebook 文档标题右上角的 Python 环境选择器来完成此操作。

设置数据集创建工具

训练自定义 YOLOv5 或任何 AI 模型的最重要部分是获得足够大且多样化的标注数据来训练模型。幸运的是，有几个数据集存储库可用。其中之一是 Google 的 Open Images。该存储库包含用于目标检测的 1600 万个边界框，涵盖 190 万张图像中的 600 个对象类别。

为了管理从此存储库中选择图像和创建数据集，Google 与 Voxel51 合作，将 Open Images 集成到其 FiftyOne 中。“FiftyOne 是一个开源工具，用于构建高质量数据集和计算机视觉模型。”该工具同时包含用于管理和查看来自各种来源（包括 Open Images）的图像数据集的 Python 包和 UI。

在附带的 Jupyter Notebook 中，您会在顶部看到以下内容

# install the FiftyOne image dataset management tool
%pip install –upgrade pip
%pip install fiftyone

设置模型训练工具

我们将使用 Ultralytics YOLOv5 GitHub Repository 提供的工具来训练和验证我们的模型。我们将在此目录的 yolov5 子目录中克隆该存储库的副本。执行以下命令

#clone YOLO v5 and install dependencies
!git clone https://github.com/ultralytics/yolov5

设置依赖项

您是否设置了 Notebook 的虚拟环境？

重要的是要为工作簿设置虚拟环境，这样您就不会将包安装到全局 Python 包中，从而可能污染依赖于全局包存储的其他系统。在 Notebook 的右上角，您可以设置 Visual Studio Code 用于从 Notebook 运行 Python 的虚拟环境。选择之前创建的虚拟环境 venv。

然后，您可以运行以下两个单元格之一，具体取决于您是否有 Nvidia GPU。

# if running on a system with GPU 

%pip install --upgrade pip 

%pip install -r requirements-gpu.txt 

%pip install ipywidgets 

# if running on a system without GPU 

%pip install --upgrade pip 

%pip install -r requirements-cpu.txt 

%pip install ipywidgets

使用小型数据集训练模型

从 Open Images 下载子集

从包含数千甚至数万张图像的大型数据集中训练模型将花费数小时甚至数天。

为了评估创建数据集和训练它的过程，我们将

• 为浣熊、狗、猫、松鼠和臭鼬各创建最多 1000 张图像的数据集
• 将数据集导出为 YOLOv5 格式
• 训练数据集并评估其性能指标。

此数据集的主要目的是检测我们友好的“垃圾熊猫”是否侵入我们的垃圾箱。我们添加了除浣熊以外的其他类别，以帮助确保生成的模型不会最终认为任何四足动物都是浣熊。

运行以下脚本会将 Open Images 数据集的子集下载到我们的 /users/<username>/.FiftyOne 目录中。此外，它还会创建多个数据子集，并将此信息存储在我们的 /users/<username>/FiftyOne 目录中的 MongoDb 数据库中。

import fiftyone as fo
import fiftyone.zoo as foz

splits = ["train", "validation", "test"]
numSamples = 1000
seed = 42

# Get 1000 images (maybe in total, maybe of each split) from fiftyone. 
# We'll ask FiftyOne to use images from the open-images-v6 dataset 
# and store information of this download in the
# dataset named "open-imges-critters". 

# The data that's downloaded will include the images, annotations, and a summary of what's
# been downloaded. That summary will be stored /Users/<username>/.FiftyOne in a mongoDB. 
# The images / annoations will be un /Users/<username>/FiftyOne.

if fo.dataset_exists("open-images-critters"):
    fo.delete_dataset("open-images-critters")

dataset = foz.load_zoo_dataset(
    "open-images-v6",
    splits=splits,
    label_types=["detections"],
    classes="Raccoon",
    max_samples=numSamples,
    seed=seed,
    shuffle=True,
    dataset_name="open-images-critters")

# Take a quick peek to see what's there
print(dataset)


# Do the same for cats, dogs, squirrels, and skunks, but after each download
# we'll merge the new downloaded dataset with the existing open-images-critters
# dataset so we can build up one large, multi-class set

if fo.dataset_exists("open-images-cats"):
    fo.delete_dataset("open-images-cats")

cats_dataset = foz.load_zoo_dataset(
    "open-images-v6",
    splits=splits,
    label_types=["detections"],
    classes="Cat",
    max_samples=numSamples,
    seed=seed,
    shuffle=True,
    dataset_name="open-images-cats")

# Now merge this new set with the existing open-images-critters set
dataset.merge_samples(cats_dataset)

if fo.dataset_exists("open-images-dogs"):
    fo.delete_dataset("open-images-dogs")

dogs_dataset = foz.load_zoo_dataset(
    "open-images-v6",
    splits=splits,
    label_types=["detections"],
    classes="Dog",
    max_samples=numSamples,
    seed=seed,
    shuffle=True,
    dataset_name="open-images-dogs")

dataset.merge_samples(dogs_dataset)

if fo.dataset_exists("open-images-squirrels"):
    fo.delete_dataset("open-images-squirrels")

squirrels_dataset = foz.load_zoo_dataset(
    "open-images-v6",
    splits=splits,
    label_types=["detections"],
    classes="Squirrel",
    max_samples=numSamples,
    seed=seed,
    shuffle=True,
    dataset_name="open-images-squirrels")

dataset.merge_samples(squirrels_dataset)

if fo.dataset_exists("open-images-skunks"):
    fo.delete_dataset("open-images-skunks")

skunks_dataset = foz.load_zoo_dataset(
    "open-images-v6",
    splits=splits,
    label_types=["detections"],
    classes="Skunk",
    max_samples=numSamples,
    seed=seed,
    shuffle=True,
    dataset_name="open-images-skunks")
    
dataset.merge_samples(skunks_dataset)

# For whenever you want to see what's been loaded.
print(fo.list_datasets())

# uncomment the following line if you wish to explore the resulting datasets in the FiftyOne UI
# session = fo.launch_app(dataset, port=5151)

将数据集导出为 YOLOv5 格式

在我们可以使用 YOLOv5 训练模型之前，我们需要将 open-images-critters 数据集导出为正确的格式。

幸运的是，FiftyOne 提供了执行此转换的工具。以下代码会将数据集导出到 datasets\critters 子文件夹。这将具有以下结构

import fiftyone as fo

export_dir = "datasets/critters"
label_field = "detections"  # for example

# The splits to export
splits = ["train", "validation","test"]

# All splits must use the same classes list
classes = ["Raccoon", "Cat", "Dog", "Squirrel", "Skunk"]

# The dataset or view to export
# We assume the dataset uses sample tags to encode the splits to export
dataset_or_view = fo.load_dataset("open-images-critters")

# Export the splits
for split in splits:
    split_view = dataset_or_view.match_tags(split)
    split_view.export(
        export_dir=export_dir,
        dataset_type=fo.types.YOLOv5Dataset,
        label_field=label_field,
        split=split,
        classes=classes,
    )

更正 dataset.yml 文件

文件 datasets\critters\dataset.yaml 在此过程中创建。

names:
- Raccoon
- Cat
- Dog
- Squirrel
- Skunk
nc: 5
path: c:\Dev\YoloV5_Training\datasets\critters
test: .\images\test\
train: .\images\train\
validation: .\images\validation\

训练代码需要标签 val 而不是 validation。将文件更新为

names: 
- Raccoon 
- Cat 
- Dog 
- Squirrel 
- Skunk 
nc: 5 
path: c:\Dev\YoloV5_Training\datasets\critters 
test: .\images\test\ 
train: .\images\train\ 
val: .\images\validation\

现在我们准备好训练模型了。

训练小型数据集

为了确保我们的流程正确，我们将训练一个具有少量 epoch（迭代次数）的模型。我们将使用标准的 yolov5s.pt 模型作为起点，训练 50 个 epoch。训练结果将存储在 train/critters/epochs50 中。在此目录中，您还将找到生成的权重（模型）以及详细介绍过程和生成性能指标的图表和表格。

由于内存限制，我们将批量大小减少到 32。请确保您已关闭任何占用大量内存的应用程序，例如 Docker，否则您的训练可能会无声地停止。

运行以下脚本进行训练。这在我的机器上花费了 51 分钟。

!python yolov5/train.py --batch 32 --weights
yolov5s.pt --data datasets/critters/dataset.yaml --project train/critters
--name epochs50 --epochs 50

完成后，train/critters.epochs50 目录将包含一些有趣的文件。

• results.csv 包含每个 epoch 的性能信息
• 几个带有各种性能指标图的 PNG 文件。PR_curve.png 和 results.png 特别值得关注。
  • PR_Curve.png 显示了最佳模型的 Precision-Recall 曲线。
  • results.png 显示了每个训练 epoch 的训练指标值。

Precision/Recall 曲线

训练指标

生成的 best.pt 的 mAP@50 为 0.757，这相当不错，但较低的 mAP@[0.5:0.95] 0.55 表明它会漏掉一些对象。

因此，让我们验证模型 train\critters\epochs50\weights\best.pt 实际上可以检测到图像中的浣熊。我将使用图像 datasets\critters\images\validation\8fbdeff053852ee7.jpg 进行此操作。您可能希望使用不同的图像。

为了了解性能，请运行两次检测。第一次运行有一些设置，导致推理时间较长，因此第二次推理时间将更具代表性。

import torch
model2 = torch.hub.load('ultralytics/yolov5', 'custom', 'train/critters/epochs50/weights/best.pt', device="0")
#print(model)
result = model2("datasets/critters/images/validation/8fbdeff053852ee7.jpg")
result.print()

Using cache found in C:\Users\matth/.cache\torch\hub\ultralytics_yolov5_master
YOLOv5  2022-11-8 Python-3.9.13 torch-1.13.0+cu117 CUDA:0 (NVIDIA GeForce RTX 3060, 12288MiB)

Fusing layers... 
Model summary: 157 layers, 7023610 parameters, 0 gradients, 15.8 GFLOPs
Adding AutoShape... 
image 1/1: 768x1024 1 Raccoon
Speed: 13.0ms pre-process, 71.0ms inference, 4.0ms NMS per image at shape (1, 3, 480, 640)

改进模型

考虑到标准的 yolov5s.pt 模型的 m_AP@50 为 0.568，我们模型的 0.757 似乎令人印象深刻。然而，之前的此过程运行表明，生成的模型会漏掉一些对象，特别是如果它们很小。

根据 Ultralytics YOLOv5 文档，改进模型性能应通过两种方式进行

• 将 epoch 数量增加到至少 300
• 增加训练集中的图像数量

我们将分两个阶段进行，以便我们可以看到每项更改对性能指标的影响。

首先，我们将训练 300 个 epoch。此外，我们将使用我们刚刚训练的模型作为起点。没必要浪费那些工作。

这将花费几个小时（在我的机器上需要 5 个小时），所以您可能想在晚上进行。

!python yolov5/train.py --batch 32 --weights
train/critters/epochs50/weights/best.pt --data datasets/critters/dataset.yaml
--project train/critters --name epochs300 -

完成后，train/critters/epochs300 文件夹将包含训练结果。

性能可以在 PR_curve.png 和 results.png 图表中看到。

Precision/Recall 曲线

训练指标

m_AP@50 的新值为 0.777，比之前提高了 2.6%。此改进的 mAP@[0.5:0.95] 0.62 表明模型会漏掉更少的对象。

再次，让我们验证它是否有效。

import torch
model2 = torch.hub.load('ultralytics/yolov5', 'custom', 'train/critters/epochs300/weights/best.pt', device="0")
#print(model)
result = model2("datasets/critters/images/validation/8fbdeff053852ee7.jpg")
result.print()

Using cache found in C:\Users\matth/.cache\torch\hub\ultralytics_yolov5_master
YOLOv5  2022-11-8 Python-3.9.13 torch-1.13.0+cu117 CUDA:0 (NVIDIA GeForce RTX 3060, 12288MiB)

Fusing layers... 
Model summary: 157 layers, 7023610 parameters, 0 gradients, 15.8 GFLOPs
Adding AutoShape... 
image 1/1: 768x1024 1 Raccoon
Speed: 8.0ms pre-process, 71.4ms inference, 2.0ms NMS per image at shape (1, 3, 480, 640)

使用大型数据集训练模型

在展示了我们可以训练自定义 YOLOv5 数据集并获得小型数据集的合理性能后，我们想尝试使用更大的数据集。然后我们将检查额外的工作和时间是否会带来任何额外的性能。

从 Open Images 下载子集

为此

• 创建包含最多 25,000 张浣熊、狗、猫、松鼠和臭鼬图像的数据集。
• 将数据集导出为 YOLOv5 格式
• 使用我们最后的 best.pt 训练数据集 300 个 epoch，并评估其性能指标。

与之前一样，运行以下脚本会将 Open Images 数据集的子集下载到我们的 /users/<username>/.FiftyOne 目录中。

import fiftyone as fo
import fiftyone.zoo as foz

splits = ["train", "validation", "test"]
numSamples = 25000
seed = 42

# Get 25,000 images (maybe in total, maybe of each split) from Fiftyone. We'll ask FiftyOne to
# use images from the open-images-v6 dataset. Store information on this download in the
# dataset named "open-images-critters-large". 

# The data that's downloaded will include the images, annotations, and a summary of what's
# been downloaded. That summary will be stored /Users/<username>/.FiftyOne in a mongoDB. 
# The images / annoations will be un /Users/<username>/FiftyOne.

if fo.dataset_exists("open-images-critters-large"):
    fo.delete_dataset("open-images-critters-large")

dataset = foz.load_zoo_dataset(
    "open-images-v6",
    splits=splits,
    label_types=["detections"],
    classes=["Raccoon", "Dog", "Cat", "Squirrel", "Skunk"],
    max_samples=numSamples,
    seed=seed,
    shuffle=True,
    dataset_name="open-images-critters-large")

print(fo.list_datasets())

# Take a quick peek to see what's there
print(dataset)

#session = fo.launch_app(dataset, port=5151)

导出大型数据集

与之前一样，我们将数据集导出为 YOLOv5 格式。我们将将其保存在 datasets/critters-large 中。

import fiftyone as fo

export_dir = "datasets/critters-large"
label_field = "detections"  # for example

# The splits to export
splits = ["train", "validation","test"]

# All splits must use the same classes list
classes = ["Raccoon", "Cat", "Dog", "Squirrel", "Skunk"]

# The dataset or view to export
# We assume the dataset uses sample tags to encode the splits to export
dataset_or_view = fo.load_dataset("open-images-critters-large")

# Export the splits
for split in splits:
    split_view = dataset_or_view.match_tags(split)
    split_view.export(
        export_dir=export_dir,
        dataset_type=fo.types.YOLOv5Dataset,
        label_field=label_field,
        split=split,
        classes=classes,
    )

更正 dataset.yml 文件

与之前一样，文件 datasets\critters-large\dataset.yaml 在此过程中创建，需要进行更正。

names:
- Raccoon
- Cat
- Dog
- Squirrel
- Skunk
nc: 5
path: c:\Dev\YoloV5_Training\datasets\critters-large
test: .\images\test\
train: .\images\train\
validation: .\images\validation\

训练代码需要标签 val 而不是 validation。将文件更新为

names:
- Raccoon
- Cat
- Dog
- Squirrel
- Skunk
nc: 5
path: c:\Dev\YoloV5_Training\datasets\critters-large
test: .\images\test\
train: .\images\train\
val: .\images\validation\

现在我们准备好训练模型了。

使用更大的数据集进行训练

同样，由于内存限制，我们需要减小批量大小。在这种情况下，我将其设置为 24。

!python yolov5/train.py --batch 24 --weights
train/critters/epochs300/weights/best.pt --data
datasets/critters-large/dataset.yaml --project train/critters-large --name
epochs300 --epochs 300 ^C

我不得不将运行停止了 15 小时。幸运的是，训练可以恢复。每个 epoch 之后，会在 weights 目录中创建一个 best.pt 文件，令人惊讶的是，它包含了迄今为止找到的最佳模型。

我们可以使用验证脚本获取此模型的性能指标。

!python yolov5/val.py --weights
train/critters-large/epochs300/weights/best.pt --data
datasets/critters-large/dataset.yaml --project val/critters --name large
--device 0

验证完成后，PR 曲线和其他图表将在 val/critters/large 文件夹中提供。

我们已经可以看到，该模型比之前的模型有了很大的改进，m_AP@50 为 0.811，比之前提高了 4.3%。

然后我可以使用训练脚本上的 --resume 参数重新启动训练

!python yolov5/train.py --resume
train/critters-large/epochs300/weights/last.pt

完成后，train/critters-large/epochs300 文件夹将包含训练结果。

性能可以在 PR_curve.png 图表中看到。

Precision/Recall 曲线

训练指标

mAP@50 的新值为 0.878，比 300 epoch 的小型模型提高了 13%，比 50 epoch 的小型模型提高了 16%。此外，mAP@[0.5:0.95] 0.75 表明模型将检测到大多数对象。

此外，mAP 图表在训练结束时的上升趋势表明，额外的训练可能会进一步改进模型。

再次，让我们验证它是否有效。

import torch
model2 = torch.hub.load('ultralytics/yolov5', 'custom', 'train/critters-large/epochs300/weights/best.pt', device="0")
#print(model)
result = model2("datasets/critters/images/validation/8fbdeff053852ee7.jpg")
result.print()

验证模型

YOLOv5 代码提供了验证自定义模型性能的工具。我们将分别在有和没有图像增强的情况下进行模型验证。

增强会操纵被推断的图像，从而向模型提供具有不同修改的多个图像。当然，这意味着任何额外的准确性都需要更多的处理时间。

至少运行两次以下脚本，因为第一次推理有一些设置。

!python yolov5/val.py --weights
train/critters-large/epochs300/weights/best.pt --data
datasets/critters-large/dataset.yaml --project val/critters-large --name
augmentee --device 0 --augment

!python yolov5/val.py --weights
train/critters-large/epochs300/weights/best.pt --data
datasets/critters-large/dataset.yaml --project val/critters-large --name not

这些验证运行的结果显示在 PR_curve.png 图表中。

正常推理 Precision/Recall 曲线

增强推理 Precision/Recall 曲线

• 未增强的验证具有 mAP@50 为 0.877，mAP@50:95 为 0.756，耗时 6.4ms。
• 增强的验证具有 mAP@50 为 0.85，mAP@50:95 为 0.734，耗时 14.0ms。

这表明增强推理并未提高性能，实际上略微降低了性能并且耗时加倍。因此，至少对于此模型，不建议进行增强推理。

结论

我们已经表明，使用现成的工具和数据，可以相对简单（尽管耗时）地

• 创建一个适合 YOLOv5 训练的大型标注数据集
• 训练并验证自定义 YOLOv5 模型的操作

虽然我没有在这里完成整个过程，但我尝试了使用

• 冻结除最后一层之外的所有层
• 冻结骨干网络（前 10 层）
• 超参数进化

前两种方法没有带来任何显著改进，或者实际上性能有所下降。

最后一种方法耗时过长而被终止。这留待以后再进行。

后续步骤

对于您，请训练一个自定义模型来检测您认为有用的内容。这可能包括检测

• 您前门上的亚马逊包裹
• 您的喂鸟器上的鸟类及其物种
• 校车来接送您的孩子
• 枪支、刀具和其他武器
• 人脸检测
• 面罩检测或未佩戴面罩检测

可能性是无限的。

对我而言，我希望

• 评估超参数的修改，以提供更好、更快的最佳模型收敛。
• 评估训练过程中超参数的演变。
• 由于 YOLOv5 现在支持分类和分割，因此可以研究这些类型模型的训练和使用
• 研究将数据集创建和训练过程封装到一个简单的工具/UI 中。

请留意后续文章。

一如既往，欢迎评论、错误报告和任何改进建议。

关于训练指标的说明

这些内容在 目标检测/实例分割的 AP 和 mAP 指标的困惑 和 平均精度 (mAP) 详解：您需要知道的一切 中有详细解释。我在这里总结了一下

其中

• Tp 是正确正面结果的数量
• Fp 是不正确的正面结果数量
• Tn 是正确负面结果的数量
• Fn 是不正确的负面结果数量。
• Ai 是两个边界框的交集面积
• Au 是两个边界框的并集面积
• APi 是类别 i 的 AP
• nc 是类别数