AI Education For All, From Your First Steps to Mastery!
Custom Object Detection Training using YOLOv5 | LearnOpenCV
使用 YOLOv5 进行自定义目标检测训练
Table of Contents
深度学习领域在 2012 年开始起飞。大约在那个时候&＃xff0c;它是一个有点排他性的领域。 我们看到编写深度学习程序和软件的人要么是深度学习从业者&＃xff0c;要么是在该领域拥有丰富经验的研究人员&＃xff0c;要么是具有非常好的编码技能的人。
今天&＃xff0c;仅仅过了10年左右&＃xff0c;情况就发生了翻天覆地的变化&＃xff0c;而且变得更好了。 从字面上看&＃xff0c;一个只学习了几周深度学习的学生可以在 20 行代码内训练一个神经网络模型。 不仅仅是针对基准数据集的现成培训。 我们正在讨论使用一些最好的模型对自定义数据集进行训练。 不相信吗&＃xff1f; 好吧&＃xff0c;这正是我们将在这篇文章中使用 YOLOV5 进行自定义对象检测训练的内容。
我们将在这篇博文中介绍以下几点&＃xff1a;
我们将在自定义数据集上训练 YOLOv5s&＃xff08;small&＃xff09;和 YOLOv5m&＃xff08;medium&＃xff09;模型。
我们还将检查冻结模型的某些层如何导致每个 epoch 的迭代时间更快&＃xff0c;以及它会对最终结果产生什么影响。
除此之外&＃xff0c;我们将比较模型的性能&＃xff0c;包括 mAP、FPS 以及 CPU 和 GPU 上的推理时间。
YOLOv5 是一种流行的实时目标检测器。 它是 YOLO&＃xff08;You Only Look Once&＃xff09;单次检测器的 PyTorch 实现&＃xff0c;它以其极快的速度和合理的准确性而闻名。
官方称&＃xff0c;作为 Darknet 框架的一部分&＃xff0c;YOLO 有四个版本。 从 YOLOv1 到 YOLOv4。 Darknet 框架是用 C 和 CUDA 编写的。
YOLOv5 是 YOLO 系列中的下一个等效版本&＃xff0c;但有一些例外。
该项目由 Glenn Jocher 在 GitHub 上的 Ultralytics 组织下启动Ultralytics organization on GitHub.。
它是使用 Python 语言编写的&＃xff0c;使用的框架是 PyTorch。
它本身就是对象检测模型的集合。 从能够在边缘设备上提供实时 FPS 的非常小的模型开始&＃xff0c;到用于云 GPU 部署的非常大且准确的模型。 它几乎拥有人们可能需要的一切。
它具有许多其他特性和功能&＃xff0c;使其成为当今任何人甚至想到对象检测时首选的对象检测模型/存储库的选择。 我们将很快检查出来。
要开始探索 YOLOv5 的整个景观&＃xff0c;让我们从模型开始。 它总共包含5个模型。 从 YOLOv5 nano&＃xff08;smallest and fastest&＃xff09;到 YOLOv5 extra-large&＃xff08;the largest model&＃xff09;。
以下是对其中每一项的简短描述&＃xff1a;
YOLOv5n&＃xff1a;它是新推出的 nano 模型&＃xff0c;是家族中最小的模型&＃xff0c;适用于边缘、物联网设备&＃xff0c;并且还支持 OpenCV DNN。 INT8 格式小于 2.5 MB&＃xff0c;FP32 格式约为 4 MB。 它是移动解决方案的理想选择。
YOLOv5s&＃xff1a;它是家族中的小型模型&＃xff0c;大约有 720 万个参数&＃xff0c;非常适合在 CPU 上运行推理。
YOLOv5m&＃xff1a;这是中型模型&＃xff0c;有 2120 万个参数。 它可能是最适合大量数据集和训练的模型&＃xff0c;因为它在速度和准确性之间提供了良好的平衡。
YOLOv5l&＃xff1a;是YOLOv5家族的大型模型&＃xff0c;4650万个参数。 它非常适合我们需要检测较小对象的数据集。
YOLOv5x&＃xff1a;是五款中最大的&＃xff0c;也是五款中mAP最高的。 虽然它比其他的慢并且有 8670 万个参数。
下图更好地概述了所有模型&＃xff0c;包括 CPU、GPU 上的推理速度&＃xff0c;以及图像大小为 640 的参数数量。
所有模型检查点都可以从存储库下载。 他们都已经在 MS COCO 数据集上进行了 300 个 epoch 的预训练。
如果您浏览存储库&＃xff0c;很明显它使自定义数据集的训练和推理变得非常容易。 可以肯定地说&＃xff0c;如果您已经准备好格式正确的数据集&＃xff0c;您可以在 2 分钟内开始训练。
但训练和推理并不是全部。 它包含许多其他功能&＃xff0c;使其非常特别。 让我们回顾一下
我们已经在上面讨论过&＃xff0c;我们可以根据用例和数据集从 5 种不同的模型中进行选择。 无论是为边缘训练实时检测器&＃xff0c;还是在云 GPU 上部署最先进的对象检测模型&＃xff0c;它都具备人们可能需要的一切。
仅对模型进行训练和推理还不足以完成目标检测管道。 在现实生活中的用例中&＃xff0c;部署也是一项主要要求。 在部署之前&＃xff0c;我们主要需要将训练好的模型转换&＃xff08;导出&＃xff09;为正确的格式。
我们可以将原生 PyTorch (.pt) 模型转换为如下格式&＃xff1a;
TorchScript
ONNX
OpenVINO
TensorRT
CoreML
TensorFlow SavedModel、GraphDef、Lite、Edge TPU 和 TensorFlow.js
这为任何深度学习工程师打开了无数的部署选项。
YOLOv5 存储库默认提供 TensorBoard 和 Weights&Biases 日志记录。 尽管您可能需要创建一个 Weights&Biases 帐户并提供 API 凭据&＃xff0c;然后才能开始训练以进行正确的日志记录。
即使您跳过它&＃xff0c;TensorBoard 日志也已经存在&＃xff0c;其中包含每个指标、损失和图像以及所有验证预测。 这使我们更容易在训练模型后随时查看性能指标。
虽然许多对象检测模型根据 MS COCO 数据集使用预定义的锚框&＃xff0c;但 YOLOv5 采用了不同的方法。 事实上&＃xff0c;之前的 YOLO 版本&＃xff0c;比如 YOLOv2&＃xff0c;只使用了 k-Means 聚类。
但是 YOLOv5 使用遗传算法来生成锚框。 他们将此过程称为自动锚定&＃xff0c;如果默认锚框不好&＃xff0c;它会重新计算锚框以适应数据。 这与 k-Means 算法结合使用以创建 k-Means 进化锚框。 这就是为什么 YOLOv5 即使在不同的数据集上也能如此出色地工作的原因之一。
YOLOv5 模型训练和检测结果如此出色的另一个原因是马赛克增强。
说明&＃xff1a;马赛克增强的示例&＃xff08;图像源image source&＃xff09;。
简单来说&＃xff0c;就是将 4 张不同的图像合二为一&＃xff0c;这样模型就可以学会处理各种复杂的图像。 它还使用其他增强技术以及马赛克增强。
现在我们已经了解了 YOLOv5 中可用的不同模型以及代码库提供的不同功能&＃xff0c;让我们将 YOLOv5 与之前的 YOLOv3 模型进行比较。 这些比较主要集中在具有 608×608 图像尺寸的模型的 mAP 上&＃xff0c;包括 Darknet 和 Ultralytics 模型。
很明显&＃xff0c;Ultralytics YOLOv3-SPP 模型在 mAP 方面能够击败 Darknet YOLOv3 SPP 模型。 而 Ultralytics YOLOv5 模型的性能甚至更好。 约 2100 万参数模型的 YOLOv5m 能够击败 6300 万参数模型的 YOLOv3-SPP 模型。 这显示了多年来 Ultralytics 模型的改进程度。
注意&＃xff1a;在上图中&＃xff0c;你可能会发现 Darknet YOLOv3 的一些结果与原始论文相比略好一些。 原因是&＃xff0c;随着 Darknet YOLOv3 模型的更新&＃xff0c;mAP 编号也由 Ultralytics 在其存储库中更新。 这也使我们能够在所有最后更新的模型之间进行公平的比较。
在这篇博文中&＃xff0c;对于使用 YOLOv5 进行自定义对象检测训练&＃xff0c;我们将使用来自 Kaggle 的这个数据集this dataset。
该数据集包含在不同环境和照明条件下的交通车辆图像。 大多数交通图像来自印度道路。 它总共包含7个类。 它们是汽车、车牌、模糊车牌、两轮车、汽车、公共汽车和卡车。
Car, Number Plate, Blur Number Plate, Two Wheeler, Auto, Bus, and Truck.
该数据集还包含车辆和车牌的图像。 这样一个具有这些类的数据集可以构成一个很好的实时交通监控应用程序。 虽然我们不会在这篇文章中这样做&＃xff0c;但我们将完成此过程所需的第一步。 也就是说&＃xff0c;构建一个好的对象检测器。 使用 YOLOv5&＃xff0c;这将非常容易&＃xff0c;因为数据集已经是所需的格式。
该数据集包含 738 张用于训练的图像、185 张用于验证的图像和 278 张用于测试的图像。 但由于只有训练集和验证集包含真实标签&＃xff0c;我们将在本文中只使用这两个。 在继续之前&＃xff0c;这里有一些图像&＃xff0c;上面绘制了基本实况框。
让我们看看我们将在使用 YOLOv5 进行自定义训练期间涵盖的内容。
我们将从训练小型 YOLOv5 模型开始。
然后我们将训练一个medium 模型并检查与small 模型相比的改进。
接下来&＃xff0c;我们将冻结freeze 几层中等medium 模型并再次训练模型。
我们将在上述所有情况下进行推理&＃xff0c;并将 mAP 指标与视频推理期间的 FPS 进行比较。
让我们开始编码部分。 所有代码都是 Jupyter notebook 的一部分&＃xff0c;您可以从下载部分访问。
在这里&＃xff0c;我们将介绍所有必要和重要的编码部分。 这些包括&＃xff1a;
让我们回顾一下代码的所有重要部分&＃xff0c;从导入我们在笔记本中使用的模块和库开始。
import os
import glob as glob
import matplotlib.pyplot as plt
import cv2
import requests
我们将需要 glob 用于捕获目录中的文件路径&＃xff0c;matplotlib 用于可视化&＃xff0c;以及 cv2 用于读取图像。
接下来&＃xff0c;我们定义一些常量和hyperparameters超参数。
TRAIN &＃61; True
# Number of epochs to train for.
EPOCHS &＃61; 25
上面&＃xff0c;我们有一个名为 TRAIN 的布尔值。 如果这是真的&＃xff0c;那么运行代码将训练笔记本中的所有三个模型。 如果我们将值提供为 False&＃xff0c;那么如果结果目录中存在任何先前训练的模型&＃xff0c;它将用于推理。 如果我们想要进行推理&＃xff0c;这是一个很好的措施&＃xff0c;可以确保我们不需要每次都训练所有模型。
下一步是下载和准备数据集。 我们需要一个简单的辅助函数来下载数据集并进行提取。
def download_file(url, save_name):
url &＃61; url
if not os.path.exists(save_name):
file &＃61; requests.get(url)
open(save_name, &＃39;wb&＃39;).write(file.content)
else:
print(&＃39;File already present, skipping download...&＃39;)
download_file(&＃39;https://learnopencv.s3.us-west-2.amazonaws.com/traffic-vehicles-object-detection.zip&＃39;, &＃39;traffic-vehicles-object-detection.zip&＃39;)
if not os.path.exists(&＃39;Traffic Dataset&＃39;):
!unzip -q "traffic-vehicles-object-detection.zip"
else:
print(&＃39;Dataset already present&＃39;)
download_file 函数将数据集下载到当前目录&＃xff08;如果它不存在&＃xff09;并提取它。 在使用数据集进行训练之前了解数据集的目录结构也很重要。
Traffic Dataset
├── images
│ ├── test [278 entries exceeds filelimit, not opening dir]
│ ├── train [738 entries exceeds filelimit, not opening dir]
│ └── val [185 entries exceeds filelimit, not opening dir]
└── labels
├── train [738 entries exceeds filelimit, not opening dir]
├── val [185 entries exceeds filelimit, not opening dir]
我们在各自的文件夹中有图像和标签。 由于地面实况标签仅存在于 train 和 val 文件夹中&＃xff0c;我们将在训练 YOLOv5 模型时使用它们。
另一种快速准备数据集的简单方法&＃xff08;包括标记和直接导出为 YOLOv5 格式&＃xff09;是 Roboflow&＃xff0c;它是 YOLOv5 的官方数据集管理工具。
YOLOv5 训练最重要的属性之一可能是数据集 YAML 文件。 该文件包含训练和验证数据的路径以及类名。 在执行训练脚本时&＃xff0c;我们需要提供这个文件路径作为参数&＃xff0c;以便脚本可以识别图像路径、标签路径以及类名。 以下是我们这里用来训练的data.yaml文件的内容。
path: "../Traffic Dataset" # Path relative to the &＃96;train.py&＃96; script.
train: images/train
val: images/val
# Classes
nc: 7
names: [
"Car", "Number Plate", "Blur Number Plate", "Two Wheeler", "Auto", "Bus", "Truck"
]
注意&＃xff1a;以上文件中的所有路径都应该是相对于训练脚本的。 由于我们将在克隆后执行 yolov5 目录中的脚本&＃xff0c;因此&＃xff0c;我们在这里将路径的值指定为“../Traffic Dataset”。
为了使用 YOLOv5 代码库的任何功能&＃xff0c;我们需要克隆它们的存储库。 接下来的几行代码克隆存储库&＃xff0c;进入 yolov5 目录并安装运行代码可能需要的所有要求。
if not os.path.exists(&＃39;yolov5&＃39;):
!git clone https://github.com/ultralytics/yolov5.git
%cd yolov5/
!pip install -r requirements.txt
如果一切顺利&＃xff0c;您应该会看到所有要求都已成功安装
如果您还记得&＃xff0c;我们讨论了如果我们准备好数据集&＃xff0c;我们如何可以立即开始训练。 其实我们现在就是这样的状态。 如果我们只是使用 data.yaml 文件的路径执行 train.py 脚本&＃xff0c;训练将立即开始。 该脚本将选择所有可用于训练的默认参数。
但我们将采取稍微不同的方法。 我们将根据我们的要求控制训练参数。 我们将运行训练脚本三个不同的时间&＃xff0c;所有参数都略有不同。
注意&＃xff1a;所有训练和推理实验均在配备第 8 代 i7 CPU、6 GB GTX 1060 GPU 和 16 GB RAM 的机器上进行。
我们将从训练小模型的所有层开始。 这意味着&＃xff0c;尽管将加载预训练的权重&＃xff0c;但整个模型将在新数据集上进行微调。
坦率地说&＃xff0c;即使控制了所有参数&＃xff0c;它也只是一个简单的命令&＃xff0c;如下所示。
RES_DIR &＃61; set_res_dir()!python train.py --data ../data.yaml --weights yolov5s.pt \
--img 640 --epochs {EPOCHS} --batch-size 16 --name {RES_DIR}
好的&＃xff01; 让我们来看看上面代码块中发生的一切。
首先&＃xff0c;我们正在创建一个新目录来保存结果。 该路径保存在 RES_DIR 变量中。 您将在笔记本中找到 set_res_dir 的整个函数定义。 这很重要&＃xff0c;因为我们想要控制结果的保存位置。 否则&＃xff0c;脚本将在执行 train.py 脚本时创建自己的目录&＃xff0c;如 run_1、run_2 等。
现在&＃xff0c;让我们回顾一下训练脚本的所有参数。
--data&＃xff1a;这接受我们之前创建的数据集 YAML 文件的路径。在我们的例子中&＃xff0c;它是返回当前目录的一个目录&＃xff0c;这就是为什么它是 ../data.yaml。
--weights&＃xff1a;此参数接受我们要用于训练的模型。由于我们使用的是 YOLOv5 系列的小型模型&＃xff0c;因此值为 yolov5s.pt。
--img&＃xff1a;我们还可以在训练时控制图像大小。图像将在被馈送到网络之前调整为该值。我们将它们的大小调整为 640 像素&＃xff0c;这也是最常用的像素之一。
--epochs&＃xff1a;此参数用于指定 epoch 的数量。因为我们已经在上面的 EPOCHS 变量中指定了 epoch 的数量&＃xff0c;所以我们在这里提供。
--batch-size&＃xff1a;这是训练时将加载到一个批次中的样本数。虽然这里的值为 16&＃xff0c;但您可以根据可用的 GPU 内存进行更改。
--name&＃xff1a;我们可以提供一个自定义目录名称&＃xff0c;所有结果都将保存在其中。在我们的例子中&＃xff0c;我们提供了一个刚刚通过调用 set_res_dir 函数创建的路径。
执行上述代码后&＃xff0c;训练将开始&＃xff0c;这将需要一些时间&＃xff0c;具体取决于硬件。 强烈建议您在 GPU 上运行所有训练。
如果训练成功完成&＃xff0c;您将看到类似于以下内容的输出。
Epoch gpu_mem box obj cls labels img_size
0/24 3.26G 0.1054 0.09665 0.05377 54 640: 100%|███
Class Images Labels P R mAP&＃64;.5 mAP&＃64;
all 185 1980 0.907 0.0817 0.0605 0.013
...
Epoch gpu_mem box obj cls labels img_size
24/24 4.34G 0.02875 0.05944 0.006674 14 640: 100%|███
Class Images Labels P R mAP&＃64;.5 mAP&＃64;
all 185 1980 0.82 0.723 0.794 0.5
25 epochs completed in 0.202 hours.
Optimizer stripped from runs/train/results_1/weights/last.pt, 14.4MB
Optimizer stripped from runs/train/results_1/weights/best.pt, 14.4MB
Validating runs/train/results_1/weights/best.pt...
Fusing layers...
Model summary: 213 layers, 7029004 parameters, 0 gradients
Class Images Labels P R mAP&＃64;.5 mAP&＃64;
all 185 1980 0.776 0.755 0.795 0.503
Car 185 1061 0.844 0.921 0.945 0.735
Number Plate 185 174 0.8 0.799 0.786 0.409
Blur Number Plate 185 161 0.7 0.54 0.608 0.289
Two Wheeler 185 271 0.83 0.903 0.922 0.602
Auto 185 94 0.775 0.658 0.713 0.331
Bus 185 110 0.814 0.797 0.876 0.626
Truck 185 109 0.67 0.67 0.712 0.53
Results saved to runs/train/results_1
YOLOv5s 模型能够在 0.5 IoU 时实现 79.5% 的 mAP&＃xff0c;在 0.5:0.95 IoU 时达到 50.3%。
下图显示了包含当前结果目录中所有losses损失和mAP的results.png文件。
将 YOLOv5s 模型训练 25 个 epoch 后的损失和 mAP 结果。
综合考虑&＃xff0c;对于一个有大约 700 万个参数的模型&＃xff0c;这些结果一点也不差。 事实上&＃xff0c;整个训练在中档 GPU 上花费了大约 12 分钟。 此外&＃xff0c;最后&＃xff0c;您可以看到它显示了保存所有结果的自定义目录。
在训练期间&＃xff0c;代码库将每个时期each epoch的validation batches验证批次的预测保存到结果目录中。 在我们检查它们之前&＃xff0c;让我们编写一个辅助函数&＃xff0c;它将在结果目录中找到所有验证预测并显示它们。
# Function to show validation predictions saved during training.def show_valid_results(RES_DIR):
!ls runs/train/{RES_DIR}
EXP_PATH &＃61; f"runs/train/{RES_DIR}"
validation_pred_images &＃61; glob.glob(f"{EXP_PATH}/*_pred.jpg")
print(validation_pred_images)
for pred_image in validation_pred_images:
image &＃61; cv2.imread(pred_image)
plt.figure(figsize&＃61;(19, 16))
plt.imshow(image[:, :, ::-1])
plt.axis(&＃39;off&＃39;)
plt.show()
上面的函数接受结果目录路径。 我们可以将图像可视化如下。
show_valid_results(RES_DIR)
以下是其中一些结果。
由于所有注释&＃xff0c;这看起来有点混乱。 但同时也表明&＃xff0c;即使是小的 YOLOv5 模型也能取得非常好的效果。
继续推理部分。 在这里&＃xff0c;我们还需要一些辅助函数。 这些是直接使用推理脚本的函数。 因此&＃xff0c;了解它们非常重要。 此外&＃xff0c;我们将在训练中型模型时进一步重用这些功能。
# 用于推断图像的辅助函数。Helper function for inference on images.
def inference(RES_DIR, data_path):
# Directory to store inference results.
infer_dir_count &＃61; len(glob.glob(&＃39;runs/detect/*&＃39;))
print(f"Current number of inference detection directories: {infer_dir_count}")
INFER_DIR &＃61; f"inference_{infer_dir_count&＃43;1}"
print(INFER_DIR)
# Inference on images.
!python detect.py --weights runs/train/{RES_DIR}/weights/best.pt \
--source {data_path} --name {INFER_DIR}
return INFER_DIR
当提供模型保存的结果目录路径和图像/视频所在的数据路径时&＃xff0c;推理功能可用于对图像和视频进行推理。 每次这个函数都会创建一个新的推理目录来保存结果。
但这里最重要的部分是运行推理的 detect.py 脚本。 回顾这些arguments 将使我们更深入地了解它的工作原理。
--weights&＃xff1a;此参数用于提供训练权重的路径以用于推理。 如果你每次都观察得当&＃xff0c;我们使用的是最好的模型&＃xff0c;那就是 best.pt。
--source&＃xff1a;我们需要提供图片或视频所在的路径。
--name&＃xff1a;最后&＃xff0c;我们每次使用自定义目录来保存推理结果&＃xff0c;以便轻松跟踪结果。
这里要注意的一件事是&＃xff0c;我们使用相同的 detect.py 脚本来推断图像和视频。 通过检查文件扩展名并调用适当的方法&＃xff0c;相同的检测脚本可以对图像和视频进行推理。
在运行 detect.py 之后&＃xff0c;只有一个最终的辅助函数用于可视化存储在磁盘上的推理图像。 这将帮助我们直接在 Jupyter Notebook 中可视化结果。
def visualize(INFER_DIR):
# Visualize inference images.
INFER_PATH &＃61; f"runs/detect/{INFER_DIR}"
infer_images &＃61; glob.glob(f"{INFER_PATH}/*.jpg")
print(infer_images)
for pred_image in infer_images:
image &＃61; cv2.imread(pred_image)
plt.figure(figsize&＃61;(19, 16))
plt.imshow(image[:, :, ::-1])
plt.axis(&＃39;off&＃39;)
plt.show()
现在&＃xff0c;我们可以完全专注于推理并检查结果。 让我们从使用经过训练的 YOLOv5s 模型进行图像推理并将它们可视化开始。
# Inference on images.
IMAGE_INFER_DIR &＃61; inference(RES_DIR, &＃39;inference_images&＃39;)
visualize(IMAGE_INFER_DIR)
下图显示了图像推断结果。
小型模型的局限性和优势从这里非常明显。 它能够检测更大的物体&＃xff0c;如汽车和卡车。 但它错误地将一个人预测为两轮车。 尽管在最后一张图片&＃xff08;右下&＃xff09;中&＃xff0c;令人印象深刻的是&＃xff0c;它能够检测到汽车上的车牌。
接下来&＃xff0c;我们将通过执行 video_inference 函数对视频进行推理。
inference(RES_DIR, &＃39;inference_videos&＃39;)
让我们在这里查看其中一个推理结果。
Youtube视频
结果看起来不错。 但我们还没有真正的比较。 虽然我们可以看到卡车的预测在路的另一边波动很大。
在 GTX 1060 GPU 上&＃xff0c;每帧的平均时间约为 8 毫秒&＃xff0c;相当于 125 FPS。 考虑到我们在这里得到的预测&＃xff0c;FPS 相当不错。
现在是训练中型模型的时候了。 训练这也将为我们提供一个机会&＃xff0c;将模型的结果与之前的结果进行比较。
我们将使用相同的训练脚本&＃xff0c;同时仅提供新的模型名称。
RES_DIR &＃61; set_res_dir()
if TRAIN:
!python train.py --data ../data.yaml --weights yolov5m.pt \
--img 640 --epochs {EPOCHS} --batch-size 16 --name {RES_DIR}
中型模型需要大约 30 分钟来训练 25 个 epoch。 以下是最终的mAP结果。
Model summary: 290 layers, 20877180 parameters, 0 gradients
Class Images Labels P R mAP&＃64;.5 mAP&＃64;
all 185 1980 0.841 0.756 0.822 0.554
Car 185 1061 0.889 0.926 0.952 0.781
Number Plate 185 174 0.84 0.799 0.845 0.471
Blur Number Plate 185 161 0.791 0.528 0.69 0.328
Two Wheeler 185 271 0.839 0.867 0.902 0.646
Auto 185 94 0.802 0.66 0.688 0.37
Bus 185 110 0.802 0.845 0.887 0.646
Truck 185 109 0.924 0.668 0.789 0.636
Results saved to runs/train/results_2
中等模型的结果要好得多。 0.5 IoU 时的 mAP 为 82.2%&＃xff0c;明显高于小模型。
此外&＃xff0c;在这种情况下&＃xff0c;损失的减少看起来更加稳定。
现在&＃xff0c;让我们仔细看看这里的图像推理结果。
从上图可以看出推理结果的提升是相当明显的。 事实上&＃xff0c;它能够检测到小模型遗漏的汽车上的车牌。
以下是与上述相同视频的视频推理结果&＃xff0c;用于比较。 检测在这里看起来很相似。 但这次远方车辆的波动较小。 此外&＃xff0c;它现在似乎更有信心检测到遥远的车牌。
Youtube视频
即使是中型模型&＃xff0c;在 GTX 1060 GPU 上似乎也能以令人印象深刻的 62 FPS 运行&＃xff0c;并获得了更多改进的结果。 这不仅仅是具有非常好的预测的实时速度。
在目标检测中&＃xff0c;我们通常使用在 MS COCO 数据集上预训练的模型&＃xff0c;并在我们自己的数据集上对其进行微调。 大多数时候&＃xff0c;我们训练模型的所有层&＃xff0c;因为对象检测是一个特别难以解决的问题&＃xff0c;数据集变化很大。
但我们不需要总是训练整个模型。 预训练模型非常强大&＃xff0c;我们可以通过冻结一些层并训练其他层来获得几乎相同的结果。 唯一的问题是这样做不是很简单。 很多时候&＃xff0c;我们需要修改源代码来冻结图层。
但是 Ultralytics YOLOv5 使得冻结模型的几层并训练其他层变得非常容易。 我们只需要在执行 train.py 脚本时使用 --freeze 参数。
现在让我们训练模型并检查上述任何理论是否正确。
RES_DIR &＃61; set_res_dir()
if TRAIN:
!python train.py --data ../data.yaml --weights yolov5m.pt \
--img 640 --epochs {EPOCHS} --batch-size 16 --name {RES_DIR} \
--freeze 0 1 2 3 4 5 6 7 8 9 10 11 12 13 14
YOLOv5 中型模型共有 25 个区块&＃xff08;从 0 到 24&＃xff09;。 每个块都是不同层的堆叠。 当我们冻结前 15 个块时&＃xff0c;卷积和批量归一化权重被冻结。 如果您在训练开始之前查看输出&＃xff0c;您将看到类似于以下的输出。
freezing model.0.conv.weight
freezing model.0.bn.weight
freezing model.0.bn.bias
.
.
.
freezing model.14.conv.weight
freezing model.14.bn.weight
freezing model.14.bn.bias
这里&＃xff0c;model.# 表示块编号&＃xff0c;然后是层类型。 正如我们所看到的&＃xff0c;从块 0 到 14 的所有卷积和批量归一化权重都被冻结了。 只剩下 10 个可训练的块。
以下块显示了训练 25 个 epoch 后的结果。
Model summary: 290 layers, 20877180 parameters, 0 gradients
Class Images Labels P R mAP&＃64;.5 mAP&＃64;
all 185 1980 0.803 0.693 0.771 0.488
Car 185 1061 0.874 0.92 0.948 0.748
Number Plate 185 174 0.797 0.812 0.84 0.433
Blur Number Plate 185 161 0.828 0.42 0.606 0.243
Two Wheeler 185 271 0.83 0.882 0.93 0.62
Auto 185 94 0.721 0.521 0.557 0.281
Bus 185 110 0.813 0.712 0.812 0.563
Truck 185 109 0.756 0.587 0.706 0.525
Results saved to runs/train/results_3
有趣的是&＃xff0c;mAP 值低于 YOLOv5 完全训练的中等模型和 YOLOv5s 模型。 但这里要注意一件事&＃xff0c;结果还不错。
此外&＃xff0c;查看这些图像推断结果。
很明显&＃xff0c;这些结果优于小型模型&＃xff0c;但不如完全训练的中型模型。
运行视频推理将提供相同的 FPS&＃xff0c;因为在推理过程中层数不会受到影响。 这个实验告诉我们&＃xff0c;如果有必要&＃xff0c;我们可以通过冻结一些层和缩短训练时间来获得同样好的结果。
从上面的实验中&＃xff0c;我们可以立即提到冻结几层和训练模型的一些好处。
首先&＃xff0c;由于我们冻结了几层&＃xff0c;这些层没有得到训练&＃xff0c;并且反向传播不会通过这些层发生。 这意味着训练迭代时间减少。 从上述训练中也可以看出&＃xff0c;与完整的中型模型训练的 30 分钟相比&＃xff0c;只需要 15 分钟即可完成。
其次&＃xff0c;即使我们冻结了几个块&＃xff0c;中等模型仍然能够产生与完全训练的模型一样好的预测。 我们可以从上面的推理结果中看到&＃xff0c;预测明显优于小模型。
性能比较
本节包含所有型号的性能比较图&＃xff0c;包括&＃xff1a;
训练时间。
平均平均精度。
GPU 上的推理速度。
CPU 上的推理速度。
这是所有三种模型的视频平均推理速度。
使用的硬件&＃xff1a;英特尔 i7 第 8 代笔记本电脑 CPU、6 GB GTX 1060 笔记本电脑 GPU。
| Model Type | GPU Inference Speed in ms (FPS) | CPU Inference Speed in ms (FPS) |
| YOLOv5s | 8.0 ms (125 FPS) | 54 ms (18 FPS) |
| YOLOv5m Full Training | 16 ms (62 FPS) | 127 ms (8 FPS) |
| YOLOv5m Frozen Layers | 16 ms (62 FPS) | 127 ms (8 FPS) |
在这篇文章中&＃xff0c;我们使用 YOLOv5 进行了大量的训练和推理实验。 我们从使用 YOLOv5 小型模型的自定义对象检测训练和推理开始。 然后我们转向了 YOLOv5 中型模型训练以及带有几个冻结层的中型模型训练。 这篇文章让我们深入了解了 YOLOv5 代码库的工作原理以及模型之间的性能和速度差异。
鉴于这篇文章中进行的大量实验&＃xff0c;您是否意识到一件事&＃xff1f; 除了一些通用的 Python 函数&＃xff0c;我们没有编写任何深度学习代码。 这表明深度学习领域正在变得多么容易&＃xff0c;并且希望它在未来也将朝着相同的方向发展。 如果您尝试对自己的数据集进行自定义训练并发现一些有趣的东西&＃xff0c;请不要忘记在评论部分分享您的结果。
Custom Object Detection Training using YOLOv5 | LearnOpenCV
如果您喜欢这篇文章并想下载本文中使用的代码&＃xff08;C&＃43;&＃43; 和 Python&＃xff09;和示例图像&＃xff0c;请单击此处。 或者&＃xff0c;注册以接收免费的计算机视觉资源指南。 在我们的时事通讯中&＃xff0c;我们分享了用 C&＃43;&＃43;/Python 编写的 OpenCV 教程和示例&＃xff0c;以及计算机视觉和机器学习算法和新闻。
https://github.com/ultralytics
https://github.com/ultralytics/yolov5
YOLOv5 &＃x1f680; in PyTorch > ONNX > CoreML > TFLite
https://github.com/ultralytics/yolov5/releases/tag/v6.1
京公网安备 11010802041100号 | 京ICP备19059560号-4 | PHP1.CN 第一PHP社区 版权所有