65.9K
CodeProject 正在变化。 阅读更多。
Home

训练 YOLOv5 模型用于人脸口罩检测

Sergio Virahonda
starIconstarIconstarIconstarIconstarIcon

5.00/5 (1投票)

2021 年 1 月 28 日

CPOL

6分钟阅读

viewsIcon

14995

在本文中，我们将使用数据集来训练 YOLOv5 对象检测模型。

引言

在本系列文章的上一篇文章中，我们标注了一个面罩数据集。现在是这个项目最激动人心的部分——模型训练。

准备训练和验证数据

与任何其他模型一样，YOLOv5 需要一些验证数据来确定在训练期间和训练后的推理效果如何。因此，我们需要将我们的图像集分割成 train 和 val 数据集，并附带相应的 .txt 文件。通常，分割比例是 80% 用于训练，20% 用于验证，并且必须按以下方式分配：

Dataset
|
| -- images
|      | -- train
|      | -- val
|
| -- labels
|      | -- train
|      | -- val
|
| -- dataset.yaml

我手动分割了数据；我将前 3,829 张图像放入 images/train 目录，其余 955 张放入 images/val 目录。然后，我相应地将它们对应的 .txt 文件复制到 labels/train 和 labels/val。

dataset.yaml 文件告诉模型数据集的分布、有多少类别以及它们的名称。在我们的例子中，该文件看起来是这样的：

功能包括

  • 第 1 行表示训练图像集的相对路径
  • 第 2 行表示验证集的相对路径
  • 第 6 行说明数据集包含多少类别
  • 第 9 行定义每个类别的名称

每次训练 YOLOv5 模型时，您都必须手动创建此文件。幸运的是，这并不难。

在 Colab Notebook 中训练 YOLOv5 模型

我们不会将您的本地计算机推向极限，因为有多种云服务选项。让我们使用Google Colab Notebook，这是一个非常强大且易于使用的解决方案。

在继续之前，我需要提及一个重要事项。YOLOv5 模型运行在PyTorch之上，PyTorch 是一个对计算资源要求很高的机器学习框架，无法在小型设备上运行。有一个解决方法——请稍等。

YOLOv5 和其他 YOLO 版本由 Ultralytics 开发，该公司维护着一个Git 仓库，您可以在其中获取使用这些模型所需的所有文件。尽管这是官方仓库，但它仍然缺少一些重要的改进。一些重要功能将于 2021 年初包含在内。该仓库将包含将 PyTorch 自定义模型转换为 TensorFlow 和 TensorFlow Lite 兼容版本所需的功能。后者在谈论 CPU/内存受限设备上的机器学习时，被认为是最佳解决方案。

上述改进是由 Ultralytics 的一位重要贡献者开发的，并可在此仓库中获得。我们将在此项目中使用它。选择此路径的原因是，当前 PyTorch – TensorFlow Lite 的转换在 Ultralytics 的流程中没有明确定义。您需要手动将 .pt 文件转换为 .onnx，然后获取 TensorFlow 权重，最后再将其转换为 TensorFlow Lite 权重。在此过程中，任何一个环节都可能出错，导致最终模型不稳定。

好了，我们开始吧！倒杯咖啡，然后启动您的Google Colab Notebook（我的在这里：我的在这里）。

让我们开始设置 GPU 实例类型。

首先，在屏幕顶部的菜单中，选择 **Edit** > **Notebook settings** > **GPU**，并确保“Hardware accelerator”（硬件加速器）设置为 **GPU**。

接下来，选择 **Runtime** > **Change Runtime Type** > **Hardware accelerator** > **GPU**。两个下拉列表都必须选择如下：

您的 Notebook 需要一些时间来初始化 - 请耐心等待几分钟。

Notebook 启动并运行后，设置初始配置：

import torch # Keep in mind that YOLOv5 runs on top of PyTorch, so we need to import it to the notebook

from IPython.display import Image #this is to render predictions

#!git clone https://github.com/ultralytics/yolov5 # Use this if you want to keep the official Ultralytics scripts.

!git clone https://github.com/zldrobit/yolov5.git
%cd yolov5
!git checkout tf-android

上面的几行将导入基本库，克隆包含模型以及 PyTorch-TensorFlow Lite 转换所需关键文件的 YOLOv5 仓库到 /content 目录。

下一步是安装一些必需的组件并更新当前 Notebook 的 TensorFlow 版本。

!pip install -r requirements.txt
!pip install tensorflow==2.3.1 #Keep this version of TF as YOLOv5 works well with it.
print('All set. Using PyTorch version %s with %s' % (torch.__version__, torch.cuda.get_device_properties(0) if torch.cuda.is_available() else 'CPU'))

让我们将我们的面罩数据集导入 Notebook。我已将其上传到GitHub，请运行以下几行：

%cd /content
!git clone https://github.com/sergiovirahonda/FaceMaskDataset

所有文件克隆完成后，您需要将数据集的 dataset.yaml 文件移动到 /content/yolov5/data 目录。使用左侧的文件浏览器，手动将文件从 /content/FaceMaskDataset 拖到 /content/yolov5/data。如前所述，此文件包含 YOLO 在自定义数据上训练模型所需的信息。如果您想检查该文件，请运行以下命令：

%cd /content
!git clone https://github.com/sergiovirahonda/FaceMaskDataset

现在是时候在自定义数据上训练模型了。首先，导航到包含 train.py 文件的 /yolov5 目录。在实现脚本时，请牢记以下几点：

  • 您必须以像素为单位指定图像大小（在本例中为 415）。模型会将每个图像的大小调整为 416 像素，因为 415 不是预定义的 max stride 32 的倍数。
  • 指定批量大小。在此例中，我们将批量大小保持为 16。
  • 您需要指定训练的 epoch 数。epoch 数越大，置信度越高。过拟合将在某个时候发生，如果数据集相对较小，过拟合会很快发生。目前 30 个 epoch 就足够了，否则您会注意到一些过拟合。
  • 通过 --data 参数，您可以引用包含数据集配置的 .yaml 文件。
  • 您需要指定预训练模型供训练管道使用。请记住，可用的选项有 YOLOv5s (yolov5s)、YOLOv5m (yolov5m)、YOLOv5l (yolov5l) 和 YOLOv5x (yolov5x)。在我们的例子中，第一个选项是最好的，因为我们需要一个非常轻量级的模型。
  • 最后，--nosave 用于仅保存最终的检查点，而 --cache 启用管道缓存图像以减少训练时间。我想保留最佳指标，所以我不会使用 --nosave。

在训练期间，为了了解训练进展情况，请密切关注 mAP@.5 指标（平均精度均值）。如果接近 1，则表示模型取得了很好的结果！

要训练模型，请运行以下命令：

%cd /yolov5
!python train.py --img 415 --batch 16 --epochs 30 --data dataset.yaml --weights yolov5s.pt --cache

上述几行代码执行对我们数据集的基本检查，缓存图像，并覆盖其他配置。它们会输出模型架构摘要（请检查以了解模型组成部分），然后开始训练。

训练结束后，/content/yolov5/runs/train/exp/weights 目录中应保存两个文件：last.pt 和 best.pt。我们将使用 best.pt。

如果您想探索训练期间记录的指标，我建议您使用 TensorBoard，这是一个非常交互式的探索工具。

%load_ext tensorboard
%tensorboard --logdir runs

这将加载类似以下内容：

请注意，mAP 和精度在 25 个 epoch 后达到最大值。这就是为什么训练您的模型超过 30 个 epoch 没有意义。

在 Google Colab 上测试模型

现在让我们探索一下模型的置信度。我们可以绘制训练期间获得的验证批次，并检查每个标签的置信分数。

Image(filename='runs/train/exp/test_batch0_pred.jpg', width=1000)

这将绘制如下内容：

虽然 0.8 分数不是非常出色，但对于此类模型来说仍然很好。请记住，YOLOv5 以牺牲准确性来换取检测速度。

现在我们有了一个训练好的模型。要在新的、未见过图像上对其进行测试，请手动创建一个新目录（例如，/content/test_images），上传一些图像，然后运行以下代码：

%load_ext tensorboard
%tensorboard --logdir runs

您将使用 best.pt 权重和 416x416 像素的图像尺寸来实现 detect.py 脚本（遵守这一点非常重要）。结果将保存到 runs/detect/exp。要显示结果，请运行以下代码：

#Plotting the first image
Image(filename='runs/detect/exp/testimage1.jpg', width=415)

结果

看起来很有希望，您觉得呢？

下一步

在下一篇文章中，我们将使我们的模型能够被 TensorFlow Lite 理解，TensorFlow Lite 是 TensorFlow 的轻量级版本，专门为在小型设备上运行而开发。敬请关注！

© . All rights reserved.