基于 PyTorch实现YOLOv5

The First Article

前言

本文记录基于PyTorch实现Github作者ultralytics的YOLOv5-tag v5.0的迁移学习大致流程，使用自制数据集（100张图片）。如有不对之处请批评指正，欢迎提出宝贵意见和建议，若需源码或者有疑问请留言。
Email：[email protected]

实现环境

使用Anaconda的虚拟环境安装以下包和其他包参考requirements.txt
Python 3.6 +PyTorch 1.10.2+opencv 4.5.4.60 +PyCharm 2021.3.2
GPU：NAVIDIA GeForce GTX 1050 Ti（4GB）

基本流程

一般分为以下四步：数据准备、建立模型、训练模型、模型评估。

数据准备

数据集包括实际场景的 100 幅图像（从百度图片下载相应图片）。使用labelImg图像标注工具，VOC格式标注图像生成标签label的xml文件，再利用脚本将标签label的xml文件转为txt文件、将数据集划分为训练集、验证集、测试集（81：9：10）。标签有七类：danger（未佩戴头盔或帽子）, hat（帽子）, person（人）, Electric cars（电动车）, car（除电动车以外的车）, Helmet（头盔）, bicycle（自行车）。这里展示训练集图片：

建立模型

从 Github上下载YOLOv5模型( yolov5 tag v5.0传送门.）进行迁移学习，修改模型配置文件，使用中规模网络配置yolov5m.yaml和预训练模型权重yolov5m.pt。

python train.py --img 640 --batch 16 --epochs 300 --data ./dataset/data.yaml --cfg models/yolov5m.yaml --weights yolov5m.pt --device 0

训练模型

训练结果如下图所示，训练完的模型在验证集下map达到82.2%，查准率90.9%，召回率72.8%。

验证集的正确标签如下图所示。

验证集的预测标签如下图所示。

模型在测试集表现如下图所示。训练完的模型在测试集下map达到81.7%，查准率81%，召回率79.2%。

测试集的正确标签如下图所示。

模型评估

如下图所示，模型在训练集、验证集是否准确预测出边框的损失值稳定收敛；模型在训练集、验证集的边框中是否准确预测出是否有物体的损失值呈现两张不同状态：训练集的损失值伴随着波动收敛、验证集的损失值迅速收敛并逐渐发散；模型在训练集、验证集的存在物体边框中是否准确预测出物体所属类别的损失值呈现两张不同状态：训练集的损失值稳定收敛、验证集的损失值迅速收敛后并上下波动；查准率与召回率在大约190轮时下降了约0.2，之后逐渐升高；Map在200轮之后趋于稳定。
results.png：

如下图所示，PR曲线中P表示precision（查准率），R表示recall（召回率），其代表的是精准率与召回率的关系。PR曲线下围成的面积即AP，所有类别AP平均值即Map。people的AP最低，hat和bicycle的AP最高。
PR_curve.png:

图片预测

视频预测

test2s