YOLOv5复现(论文复现)
本文所涉及所有资源均在传知代码平台可获取
文章目录
- YOLOv5复现(论文复现)
- 概述
- 模型结构
- 正负样本匹配策略
- 损失计算
- 数据增强
- 使用方式
- 训练
- 测试
- 验证
- Demo
概述
YOLOv5是由Ultralytics公司于2020年6月开源的目标检测模型,具有轻量化、易用性和高性能等特点,在不同的硬件平台上提供了很好的速度和性能平衡,凭借其优秀的实时性和准确性在工业、安防、无人驾驶和许多其他领域都有广泛的应用。YOLOv5是YOLO系列中最受欢迎的工作之一,但就其整体架构而言,YOLOv5可以看作是YOLOv4的“精心调教”版,通过对YOLOv4的网络结构、优化器超参、数据预处理超参、损失函数超参等多个超参数调优,使得YOLOv5的性能要远远优于YOLOv4,但整体架构仍旧延续了YOLOv4的Backbone+SPP+PaFPN+Head的结构,并无较大的改动。从结构上来看,YOLOv5仍采用了YOLOv4的CSPDarkNet结构,设计了width因子和depth因子来对模型做缩放,从而构建出了N/S/M/L/X等不同的模型尺度;从标签分配角度来看,YOLOv5依旧是anchor-based路线,且沿用了YOLOv3以来的anchor box参数,但将YOLOv4还在用的IoU-based匹配策略修改为shape-based匹配策略(长宽比阈值),可以为每个目标分配更多的正样本;从损失函数角度来看,还是objectness+classification+regression三部分损失,且损失函数与YOLOv4是一致的。从使用角度来看,YOLOv5相比YOLOv4有了很大的进步,采用了更加流行的pytorch框架,使得上手难度大大降低,并提供了更加完善的文档和更加丰富的预训练权重
模型结构
YOLOv5的结构可以大体上分为“主干网络backbone+SPP颈部网络+PaFPN特征金字塔+检测头”,如下图所示,图的上半部分为模型总览;下半部分为具体网络结构
首先是主干网络,和YOLOv4 的主干网络相同,都采用了基于CSP结构的CSPDarkNet网络,CSPNet是一种增强CNN学习能力的跨阶段局部网络。CSPNet全称是Cross Stage Partial Network,主要从网络结构设计的角度来解决以往工作在推理过程中需要很大计算量的问题,其能够在降低20%计算量的情况下保持甚至提高CNN的能力。不同的是,在YOLOv4中,CSPDarkNet遵循“12884”的设计来得到CSPDarkNet-53,而YOLOv5则进一步将其修改为“3993”的结构,以便后续在模型深度的层面上做缩放。相较于YOLOv4 的“12884”,YOLOv5则在模型的第一阶段中不加入CSP模块,仅仅使用一个kernel较大的卷积来做第一次的降采样,随后才使用CSP模块。整体上来看,和YOLOv4的CSPDarkNet-53的框架是一样的,没有本质区别,不过是在深度和宽度上做了一些调整,且第一阶段的结构不大相同:YOLOv4的第一阶段使用了CSP模块(深度为1),而YOLOv5的第一阶段仅使用一个大核卷积。为了调控YOLOv5的模型规模,一共设计了5套尺度因子,如下表所示,其中 depth 控制网络结构深度,主要作用在CSP模块中来调整其中的残差块的数量,进而调整了网络的深度;width 控制网络结构宽度,即模块输出特征图的通道数
模型尺度 | 宽度因子width | 深度因子depth |
---|---|---|
N | 0.25 | 0.34 |
S | 0.50 | 0.34 |
M | 0.75 | 0.67 |
L | 1.0 | 1.0 |
X | 1.25 | 1.34 |
CSPDarkNet实现代码如下
# CSPDarkNet
class CSPDarkNet(nn.Module):
def __init__(self, depth=1.0, width=1.0, act_type='silu', norm_type='BN', depthwise=False):
super(CSPDarkNet, self).__init__()
self.feat_dims = [round(64 * width), round(128 * width), round(256 * width), round(512 * width), round(1024 * width)]
# P1/2
self.layer_1 = Conv(3, self.feat_dims[0], k=6, p=2, s=2, act_type=act_type, norm_type=norm_type, depthwise=depthwise)
# P2/4
self.layer_2 = nn.Sequential(
Conv(self.feat_dims[0], self.feat_dims[1], k=3, p=1, s=2, act_type=act_type, norm_type=norm_type, depthwise=depthwise),
CSPBlock(in_dim = self.feat_dims[1],
out_dim = self.feat_dims[1],
expand_ratio = 0.5,
nblocks = round(3*depth),
shortcut = True,
act_type = act_type,
norm_type = norm_type,
depthwise = depthwise)
)
# P3/8
self.layer_3 = nn.Sequential(
Conv(self.feat_dims[1], self.feat_dims[2], k=3, p=1, s=2, act_type=act_type, norm_type=norm_type, depthwise=depthwise),
CSPBlock(in_dim = self.feat_dims[2],
out_dim = self.feat_dims[2],
expand_ratio = 0.5,
nblocks = round(9*depth),
shortcut = True,
act_type = act_type,
norm_type = norm_type,
depthwise = depthwise)
)
# P4/16
self.layer_4 = nn.Sequential(
Conv(self.feat_dims[2], self.feat_dims[3], k=3, p=1, s=2, act_type=act_type, norm_type=norm_type, depthwise=depthwise),
CSPBlock(in_dim = self.feat_dims[3],
out_dim = self.feat_dims[3],
expand_ratio = 0.5,
nblocks = round(9*depth),
shortcut = True,
act_type = act_type,
norm_type = norm_type,
depthwise = depthwise)
)
# P5/32
self.layer_5 = nn.Sequential(
Conv(self.feat_dims[3], self.feat_dims[4], k=3, p=1, s=2, act_type=act_type, norm_type=norm_type, depthwise=depthwise),
SPPF(self.feat_dims[4], self.feat_dims[4], expand_ratio=0.5),
CSPBlock(in_dim = self.feat_dims[4],
out_dim = self.feat_dims[4],
expand_ratio = 0.5,
nblocks = round(3*depth),
shortcut = True,
act_type = act_type,
norm_type = norm_type,
depthwise = depthwise)
)
def forward(self, x):
c1 = self.layer_1(x)
c2 = self.layer_2(c1)
c3 = self.layer_3(c2)
c4 = self.layer_4(c3)
c5 = self.layer_5(c4)
outputs = [c3, c4, c5]
return outputs
对于颈部网络,yolov5使用了和YOLOv4中类似的SPP模块。SPP(Spatial Pyramid Pooling)利用不同的池化核尺寸提取特征的方式可以获得丰富的特征信息,有利于提高网络的识别精度。对每个特征图,使用三种不同尺寸的池化核进行最大池化,分别得到预设的特征图尺寸,最后将所有特征图展开为特征向量并融合,确保输入预定义全连接层的feature vector(特征向量)是固定尺寸。
对于特征金字塔,yolov5使用了PaFPN结构,将FPN(Feature pyramid network)和PANet(Path Aggregation Network)结合起来,通过Bottom-Up和Top-down结构充分融合高层特征和底层特征。相较于YOLOv4,YOLOv5在PaFPN中添加了CSP模块,取代了早期的包含5层卷积的简单模块,同时也加入了depth因子来调整PaFPN的深度。
对于检测头,输出objectness+classification+regression。采用的是解耦检测头,将检测头的类别特征分支的输出去做classification,将检测头的位置特征分支的输出去做regression和objectness
正负样本匹配策略
正负样本匹配策略的核心是确定预测特征图的所有位置中哪些位置应该是正样本,哪些是负样本,甚至有些是忽略样本。 匹配策略是目标检测算法的核心,一个好的匹配策略可以显著提升算法性能。采用了 anchor 和 gt_bbox 的 shape 匹配度作为划分规则,同时引入跨邻域网格策略来增加正样本。对于任何一个输出层,抛弃了常用的基于 IoU 匹配的规则,而是直接采用 shape 规则匹配,也就是该 GT Bbox 和当前层的 Anchor 计算宽高比,如果宽高比例大于设定阈值,则说明该 GT Bbox 和 Anchor 匹配度不够,将该 GT Bbox 暂时丢掉,在该层预测中该GT Bbox 对应的网格内的预测位置认为是负样本。接下来, 只需要确定这些anchor box都是来自于哪个特征金字塔等级,从而去计算目标框在相应的特征金字塔等级上的中心点坐标
损失计算
YOLOv5 中总共包含 3 个 Loss,分别为:
- classification loss:使用的是 BCE loss
- Objectness loss:使用的是 BCE loss
- Regression loss:使用的是 CIoU loss
三个 loss 按照一定比例汇总,对Objectness 和classification 分别设置为1.0,对于Regression 则设置为5.0
数据增强
YOLOv5 中使用的数据增强比较多,包括:
- Mosaic 马赛克
- RandomAffine 随机仿射变换
- MixUp
- 图像模糊等采用Python的 albumentations库实现的变换
- HSV 颜色空间增强
- 随机水平翻转
其中 Mosaic 数据增强概率为 1,表示一定会触发,而对于 small 和 nano 两个版本的模型不使用 MixUp,其他的 l/m/x 系列模型则采用了 0.1 的概率触发 MixUp。小模型能力有限,一般不会采用 MixUp 等强数据增强策略。当四张640×640的图像被拼接成1280×1280的马赛克图像后,会再使用随机仿射变换从中截取出新的640×640图像,作为最终用于训练的马赛克增强
使用方式
创建python虚拟环境
conda create -n yolov5 python=3.8
conda activate yolov5
安装相关依赖
pip install -r requirents.txt
准备数据集(如果需要从头训练YOLOv5模型,则需要下载数据集,如果只是使用YOLOv5进行目标检测则可以不用下载,可以使用自己的数据进行测试)官网下载COCO数据集,如下图红色框所示
下载完成并解压后后目录如下:
F:\datasets
|___COCO2017
|___annotations
|____instances_train2017.json
|____instances_val2017.json
...
|___train2017
|____000000000009.jpg
...
|___val2017
|____000000000139.jpg
...
2.清洗COCO数据集,会在annotations目录下生成instances_train2017_clean.json和instances_val2017_clean.json
cd tools/
python clean_coco.py --root path/to/coco --image_set train
python clean_coco.py --root path/to/coco --image_set val
3.检查COCO数据集
python dataset/coco.py
下面的实验都以yolov5_s模型为例,如需换成其他模型,将命令中的参数-m yolov5_s换成其他模型即可,如-m yolov5_l
训练
使用COCO数据集从头开始训练YOLOv5模型(将下面的F:\datasets\换成自己数据集路径)
python train.py --cuda -d coco --root F:\datasets\ -m yolov5_s -bs 1 --max_epoch 300 --wp_epoch 1 --eval_epoch 10 --fp16 --ema --multi_scale
测试
使用COCO2017val数据集测试训练好的模型,会依次看到检测结果的可视化图像(将yolov5_s_coco_adamw.pth换成自己模型权重的路径)
python test.py -d coco --cuda -m yolov5_s --img_size 640 --weight yolov5_s_coco_adamw.pth --root F:\datasets\ --no_multi_labels --show
验证
使用COCO2017val数据集验证训练好的模型,会看到COCO风格的AP结果输出
python eval.py -d coco --cuda -m yolov5_s --img_size 640 --weight yolov5_s_coco_adamw.pth --root F:\datasets\
Average Precision (AP) @[ IoU=0.50:0.95 | area=medium | maxDets=100 ] = 0.444
Average Precision (AP) @[ IoU=0.50:0.95 | area= large | maxDets=100 ] = 0.519
Average Recall (AR) @[ IoU=0.50:0.95 | area= all | maxDets= 1 ] = 0.324
Average Recall (AR) @[ IoU=0.50:0.95 | area= all | maxDets= 10 ] = 0.544
Average Recall (AR) @[ IoU=0.50:0.95 | area= all | maxDets=100 ] = 0.612
Average Recall (AR) @[ IoU=0.50:0.95 | area= small | maxDets=100 ] = 0.432
Average Recall (AR) @[ IoU=0.50:0.95 | area=medium | maxDets=100 ] = 0.685
Average Recall (AR) @[ IoU=0.50:0.95 | area= large | maxDets=100 ] = 0.766
ap50_95 : 0.3912800741053746
ap50 : 0.5693696831091651
Demo
使用自己的数据测试训练好的模型,需将下面的path_to_img和path_to_vid换成图片或视频的路径,将weight换成YOLOv5模型权重的路径.
图片
python demo.py --mode image --path_to_img dataset\demo\images --cuda --img_size 640 --model yolov5_s --weight yolov5_s_coco_adamw.pth --dataset coco --num_classes 80 --show
视频
python demo.py --mode video --path_to_vid dataset\demo\videos\01.mp4 --cuda --img_size 640 -m yolov5_s --weight yolov5_s_coco_adamw.pth --show --gif
文章代码资源点击附件获取