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

目录

介绍

设置您的训练环境

硬件

开发集成开发环境

设置项目文件夹

设置Python虚拟环境

用于创建数据集的设置工具

设置模型训练工具

设置依赖项

使用小数据集训练模型

从打开的图像下载子集

将数据集导出为YOLOv5格式

更正dataset.yml文件

训练小型数据集

精度/召回率曲线

训练指标

改进模型

精度/召回率曲线

训练指标

训练模型

从打开的图像下载子集

导出大型数据集

更正dataset.yml文件

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

精度/召回率曲线

训练指标

验证模型

正态推理精度/召回率曲线

增强推理精度/召回率曲线

结论

后续步骤

关于训练指标的说明

---

• 下载小动物模型 - 12.5 MB
• 下载自定义模型训练 - 110.4 KB

介绍

对象检测是人工智能的常见应用。在CodeProject.AI服务器中，我们添加了一个使用YOLOv5架构进行对象检测的模块。这是一个出色的系统，但使用“标准”YOLOv5模型意味着您只能使用默认模型中可用的80个类。若要检测此默认设置之外的对象，需要训练自己的自定义模型。

虽然YOLOv5文档将引导您完成创建新模型的过程，但设置训练类并生成良好且准确的模型并不像它们让您相信的那么简单。

本文将引导您使用以下步骤创建自定义模型来检测后院害虫：

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

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

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

设置您的训练环境

硬件

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

规范	值
CPU	第12代智能Intel® 酷睿™ i5-12400 2.50 GHz
RAM	16 GB
GPU	NVidia GeForce RTX 3060
GPU RAM	12GB
Disk	1 TB固态硬盘
OS	Windows 11

开发集成开发环境

我正在使用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模型的最重要部分是获取足够大且多样化的注释数据集来训练模型。幸运的是，有几个数据集存储库可用。其中之一是谷歌的Open Images。此存储库包含用于对象检测的16M边界框，用于600.1M图像上的9个对象类。

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

在附加的Jupyter notebook中，您将在顶部看到以下内容：

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

设置模型训练工具

我们将使用Ultralytics YOLOv5 GitHub存储库提供的工具来训练和验证我们的模型。我们将这个仓库的副本克隆到 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

使用小数据集训练模型

从打开的图像下载子集

从包含数千或数万张图像的大型数据集训练模型将需要数小时或数天。

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

• 创建一个数据集，每个数据集最多包含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而不是验证。将文件更新为

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.epohs50 目录将包含许多有趣的文件。

• results.csv包含每个epoch的性能信息
• 几个PNG文件，其中包含各种性能指标的图表。PR_curve.png和results.png特别令人感兴趣。
  • PR_Curve.png显示最佳模型的精度与召回率曲线。
  • results.png显示训练的每个epoch的训练指标值。

精度/召回率曲线

训练指标

得到的best.pt的mAP@50值为0.757，这是相当不错的，但是mAP@[0.5:0.95] 0 = 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图表中看到。

精度/召回率曲线

训练指标

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数据集并使用小数据集获得合理的性能之后，我们想在更大的数据集上尝试一下。然后，我们将检查是否有任何额外的性能值得额外的努力和时间。

从打开的图像下载子集

为此：

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

和以前一样，运行以下脚本会将开放图像数据集的子集下载到我们的/users/<username>/中。五十一个目录。

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图中看到。

精度/召回率曲线

训练指标

mAP@50的新值为0.878，比13个epoch的小模型额外提高了300%，比具有16个epoch的小模型提高了50%。此外，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图中。

正态推理精度/召回率曲线

增强推理精度/召回率曲线

• 未增强验证的mAP@50为0.877，mAP@50：95为0.756，耗时6.4毫秒。
• 增强验证的mAP@50为0.85，mAP@50：95为0.734，耗时14.0毫秒。

这表明增强推理不会提高性能，实际上会略微降低性能，并且需要两倍的时间。因此，至少对于此模型，不建议使用增强推理。

结论

我们已经证明，使用现成的工具和数据，这是一项相对简单但耗时的任务：

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

虽然我没有在这里经历这个过程，但我确实尝试了以下方法进行训练：

• 冻结除最后一层之外的所有层
• 冻结骨架（前10层）
• 超参数演进

前两个没有带来任何显着的改进，或者实际上性能下降。

最后一个只是花了太长时间并被终止。这留待今后进行。

后续步骤

对于你来说，训练一个自定义模型来检测你认为有用的内容。这可能是在检测：

• 亚马逊包裹在你家门口。
• 鸟在你的鸟喂食器中的存在和种类。
• 校车到达接送您的孩子。
• 枪支、刀具和其他武器。
• 人脸检测
• 口罩检测，或缺乏口罩检测。

可能性是无穷无尽的。

对我来说，我想：

• 评估对超参数的修改，以在最佳模型上提供更好、更快的收敛。
• 评估训练过程中体型参数的演变。
• 由于YOLOv5现在支持分类和分段，请查看这些类型模型的训练和使用
• 了解如何将数据集创建和训练过程封装到一个简单的工具/UI中。

关于训练指标的说明

这些在对象检测/实例分割的AP和mAP的混淆指标和平均平均精度（mAP）解释：您需要知道的一切中进行了详细解释。我在这里总结了这一点：

术语	公式	描述
Precision	Precision=Tp/（Tp+Fp）	阳性结果概率的测量（对象检测）是正确的。
Recall	Recall=Tp/（Tp+Fn）	衡量本应出现阳性结果的概率。
AP	Precision-Recall率曲线下的区域	平均精度
IoU	IoU =Ai/Au	交集上的并集是两个边界框重叠量的度量。当其中一个边界框是地面实况，另一个是预测值时，IoU是预测准确性的度量。值越大，预测值和真实地面值的匹配度就越高。
mAP	mAP=1/nc∑APi从i=1到nc	平均精度。每类对象的平均精度之和的平均值。此值是针对IoU阈值的特定值计算的正值。
mAP@50		mAP表示IoU阈值为.5，表示正面预测。
mAP@[0.5：0.95]		介于.5和.95之间的多个IoU值的mAP值的平均值。通常使用10个间隔或值.05。
Ground Truth		对象的边界框的真实值。这通常由手动或自动标记过程定义，该过程定义了图像中对象的类和边界框。

其中：

• Tp是正确阳性结果的数量
• Fp是错误阳性结果的数量
• Tn是正确阴性结果的数量
• Fn是不正确的阴性结果的数量。
• Ai是两个边界框的交集面积
• Au是两个边界框的并集面积
• APi是i类的AP
• NC是类数

https://www.codeproject.com/Articles/5347827/How-to-Train-a-Custom-YOLOv5-Model-to-Detect-Objec