计算机系统应用教程网站

网站首页 > 技术文章 正文

从0开始实现目标检测——基于YOLOv5

btikc 2024-09-02 16:49:18 技术文章 13 ℃ 0 评论

一. 背景

在上次我们利用YOLOv3做目标检测的任务中(参考文章:从0开始实现目标检测——实践篇),我们遗留一些问题,其中一个是利用YOLOv5做一次对比,既然是对比,也就是在数据集不变的情况下,基于YOLOv5进行训练,看看整个模型的mAP是否有变化。整个过程也分为以下4步:

  1. 为什么是YOLOv5 ?
  2. 安装体验YOLOv5,并跑通基于公开数据训练YOLOv5
  3. 基于自己的数据训练YOLOv5
  4. 计算mAP

那么,我们开始吧~

二. 为什么是YOLOv5?

基于对目标检测原理的探究(参考文章:从0开始实现目标检测——原理篇),我们知道了YOLO系列的算法思想:它不使用窗口滑动,而是直接将原始图片分割成互不重合的小方块,然后通过卷积最后生产这样大小的特征图,可以认为特征图的每个元素也是对应原始图片的一个小方块,然后用每个元素来可以预测那些中心点在该小方格内的目标。

YOLO系列的算法都是基于这个思想实现的。对比YOLOv3,YOLOv5在整个神经网络分为4个部分的改进如下:

  1. Input:数据加载使用了3种数据增强:缩放、色彩空间调整和马赛克增强。
  2. BackBone:结合了很多先进的图像识别领域的内容和算法,包括:CSPNet、Leaky ReLU和Sigmoid 激活函数。
  3. Neck:在BackBone和最后的输出层之间往往会插入一些层,这里就加入了SPP-Net、FPN+PAN结构。
  4. Prediction:输出层的锚框机制和Yolov3相同,主要改进的是训练时的损失函数GIOU Loss,加快了收敛速度。

此外,YOLOv5还增加了自适应锚定框的功能,这样就不用根据不同训练数据调整锚定框的大小和位置了。

在这些改进下,YOLOv5在兼顾mAP的同时,有着更短的检测时间,同时YOLOv5s的权重文件大小只有27MB,能够更好的适应嵌入式设备和移动设备。这里有张官方的性能图参考,如下:

简单解释下YOLOv5上新增的技术点中的2个:

SPP-Net:神经网络是需要输入固定尺寸的图片,比如224x224、32x32、96x96等。这样对于我们希望检测各种大小的图片的时候,需要经过crop,或者warp等一系列操作,这都在一定程度上导致图片信息的丢失和变形,限制了识别精确度。而且,从生理学角度出发,人眼看到一个图片时,大脑会首先认为这是一个整体,而不会进行crop和warp,所以更有可能的是,我们的大脑通过搜集一些浅层的信息,在更深层才识别出这些任意形状的目标。SPP-Net对这些网络中存在的缺点进行了改进,基本思想是,输入整张图像,提取出整张图像的特征图,然后利用空间关系从整张图像的特征图中,在池化层中提取各个区域的特征。一个正常的深度网络由两部分组成,卷积部分和全连接部分,要求输入图像需要固定size的原因并不是卷积部分而是全连接部分。所以SPP层就作用在最后一层卷积之后,SPP层的输出就是固定大小。SPP-Net不仅允许测试的时候输入不同大小的图片,训练的时候也允许输入不同大小的图片,通过不同尺度的图片同时可以防止过拟合。训练速度也会加快。

FPN+PAN结构:FPN是一个自顶向下的,将高层特征逐渐传递下去的网络结构,类似一个金字塔,会增加模型对不同缩放尺寸物体的检测,从而识别不同大小和尺寸的同一个物体。YOLOv3就是这种网络结构的实现。而PAN是一个从下而上的FPN网络,改善了低层特征的传播。第三条通路的每个阶段都将前一阶段的特征映射作为输入,并用3x3卷积层处理它们。输出通过横向连接被添加到自上而下通路的同一阶段特征图中,这些特征图为下一阶段提供信息。同时使用自适应特征池化恢复每个候选区域和所有特征层次之间被破坏的信息路径,聚合每个特征层次上的每个候选区域,避免被任意分配。

YOLOv5的改进了解到这里了,我们继续看看如何在YOLOv5上训练数据。

三. 安装YOLOv5,并跑通COCO数据集的测试

YOLOv5是完全基于Pytorch的,而Pytorch的主要性能发挥又依赖于CUDA和cuDNN,所以一定要安装好CUDA,cuDNN,和对应版本的Pytorch,这点非常重要,要是Pytorch和CUDA,cuDNN的版本不对应,YOLOv5的训练是跑不起来的。对应关系可以从Pytorch的官网: https://pytorch.org/ 上找到。

好了,现在我们开始按照YOLOv5吧:

git clone https://github.com/ultralytics/yolov5.git
cd yolov5
pip install -r requirements.txt

安装完毕后,我们可以简单检测下,利用如下命令检测test.jpg的目标。

python detect.py --source data/images/zidane.jpg --weights yolov5s.pt

我们使用YOLOv5自带的齐达内的图片做个体验,--weights参数指定使用哪个权重文件进行推理。

命令运行后,会在项目的run目录下建立一个exp的子目录,里边保存了对图片的目标检查结果,大家可以打开看看。如图:

也可以通过摄像头进行目标检测,命令如下:

python detect.py --source 0

source参数指定了检测输入源,可以包括Video, RTMPStream, Images, Webcam等多种源。

weights参数指定了使用的权重文件,根据模型的规模不同设有4个模型:v5s、v5m、v5l、v5x,权重文件的比较如下表:

Model

APval

APtest

AP50

SpeedGPU

FPSGPU

params

FLOPS

YOLOv5s

37.0

37.0

56.2

2.4ms

416

7.5M

13.2B

YOLOv5m

44.3

44.3

63.2

3.4ms

294

21.8M

39.4B

YOLOv5l

47.7

47.7

66.5

4.4ms

227

47.8M

88.1B

YOLOv5x

49.2

49.2

67.7

6.9ms

145

89.0M

166.4B

YOLOv5x + TTA

50.8

50.8

68.9

25.5ms

39

89.0M

354.3B

YOLOv3-SPP

45.6

45.5

65.2

4.5ms

222

63.0M

118.0B

可以根据实际情况选择合适的权重文件,我们选择YOLOv5l做为我们的权重文件进行训练。

我们的思路和上次YOLOv3的一样,还是先用Coco训练集跑通训练流程,然后再替换成我们自己的训练集。运行命令进行测试:

python train.py --img 640 --batch 16 --epochs 5 --data ./data/coco128.yaml --cfg ./models/yolov5s.yaml --weights yolov5s.yaml

程序会自动下载coco128的数据集,注意这里的epochs数据并不大,主要还是为了跑通训练流程,查找环境和程序错误用。

我在训练的时候遇到的一个问题就是pytorch的版本和cuda的版本不匹配,这样在pytorch的torch.cuda.is_avaliabe()依然会返回True,但是训练的时候torch会报关员cuda的错误,最后在pytorch的官网按照cuda的版本下载安装才解决了这个问题。


训练结束后,会在runs/目录下建立一个train的文件夹保存权重文件。使用命令

python detect.py --weigths runs/train/exp/weights/best.pt --source data/images/zidane.jpg

进行图片检测测试。


同样,在runs/train/exp目录下会看到模型的mAP、数据分类数量、各个类的精度、召回率等指标图片。这点比YOLOv3要方便的多,当时YOLOv3是通过自己编写的一个脚本程序计算的。

四. 训练自己的模型

Coco128的数据集训练完成了,我们接着考虑如何训练自己的数据。先找到数据集定义的地方,看看数据集要怎么组织?

打开data/coco128.yaml的文件,内容如下:

path: ../datasets/coco128 # dataset root dir
train: images/train2017 
val: images/train2017
test:  # test images (optional)

# Classes
nc: 80  # number of classes
names: ['person', 'bicycle', 'car', 'motorcycle', 'airplane', 'bus', 'train', 'truck', 'boat', 'traffic light',
        'fire hydrant', 'stop sign', 'parking meter', 'bench', 'bird', 'cat', 'dog', 'horse', 'sheep', 'cow',
				'elephant', 'bear', 'zebra', 'giraffe', 'backpack', 'umbrella', 'handbag', 'tie', 'suitcase', 'frisbee',
				'skis', 'snowboard', 'sports ball', 'kite', 'baseball bat', 'baseball glove', 'skateboard', 'surfboard',
				'tennis racket', 'bottle', 'wine glass', 'cup', 'fork', 'knife', 'spoon', 'bowl', 'banana', 'apple',
				'sandwich', 'orange', 'broccoli', 'carrot', 'hot dog', 'pizza', 'donut', 'cake', 'chair', 'couch',
				'potted plant', 'bed', 'dining table', 'toilet', 'tv', 'laptop', 'mouse', 'remote', 'keyboard', 'cell phone',
				'microwave', 'oven', 'toaster', 'sink', 'refrigerator', 'book', 'clock', 'vase', 'scissors', 'teddy bear',
				'hair drier', 'toothbrush']  # class names
          
# Download script/URL (optional)
download: https://ultralytics.com/assets/coco128.zip
  • path: 是数据集合存放的地方,我们打开数据集的位置,就能看到组织方式。
  • train:是训练集的位置,在path的images/train2017子目录下。
  • val: 是验证集的位置,在path的images/train2017子目录下。
  • nc:是类别的数量。
  • names: 是类别名称列表。

由于上次我们在基于YOLOv3进行训练的时候,已经对数据集做了标注,详情可查看文章(参考文章:从0开始实现目标检测——实践篇),所以这次训练同样的数据集,我们就不再重复这个过程,只是把相关的数据集复制到我们自己的数据集目录下即可。

然后先复制了一份coco128.yaml文件,命名为my_data.yaml,接着把nc换成我们的要识别数量6,names是我们要识别的列表名称。

修改完了数据集,我们看看模型文件都有什么定义,根据我的情况,我打算采用v5l的模型,于是打开文件models/yolov5l.yaml,发现只需要把nc修改成和数据文件定义的一样,其他的模型参数可以暂时先不修改,训练完毕后看效果再进行调整。同样复制一份yolov5l.yaml命名为my_model.yaml,然后修改nc=6。


一切准备妥当,让我们开始训练吧~

python train.py --epochs 3000 --data data/my_data.yaml --cfg models/my_model.yaml --weights yolov5l.yaml

开始训练后,需要提供一个wandb的账号,以供训练课程可视化。这个wandb也顺带解决了上次进行YOLOv3训练时遗留的训练可视化的问题。wandb的注册和使用并不是很麻烦,这里就不再多介绍了。


第二天发现训练完毕了,这个有点超出我的预计,当时在基于YOLOv3进行训练的时候可用了4天,现在数据集、机器配置都没有改变,训练居然只用了1天,还是非常不错的。然后在runs/train/exp2下,保存了训练结果和模型的性能指标图。如下显示,mAP_0.5达到了0.83,mAP_0.5:0.95也能够到0.6,比上次在YOLOv3上的0.49提高了55%,数据结果还是非常喜人的。

我们利用我们训练的模型做了点些检查看看实际效果吧:

python detect.py --weigths runs/train/exp2/weights/best.pt --source data/images/xxx.jpg

结果如下:

五. 总结

  1. YOLOv5的训练时间、mAP都比以前YOLOv3下提高了不少。
  2. 本次的问题主要是pytorch要按照cuda和cuDNN的版本匹配安装。
  3. 本次除了基于YOLOv5完成了自有模型训练外,还解决了上次遗留的训练可视化问题。
  4. 跟上次一样,同样有个TODO List留下来,就是想把自己训练的模型装在移动设备上,然后拿到车里和路上,看看实际效果怎么样。

好了,本次的YOLOv5的训练和体验就到此结束了,感谢大家~






Tags:

本文暂时没有评论,来添加一个吧(●'◡'●)

欢迎 发表评论:

最近发表
标签列表