YOLOV5 的小目标检测网络结构优化方法汇总(附代码)

点击上方“3D视觉工坊”,选择“星标”

干货第一时间送达

18a90f31b507c92904831e0f5d45c184.png

作者丨南山

来源丨 AI约读社

YOLOv5是一种非常受欢迎的单阶段目标检测,以其性能和速度著称,其结构清晰灵活。虽然 yolov5是一个强有力的工具,但它被设计成一个通用的目标检测器,因此对较小的目标检测没有很好的优化。本文将总结了当前针对yolov5小目标检测的网络结构优化方法。 主要有以下几个方法:

1、增加小目标检测层

2、Transformer Prediction Heads (TPH)集成到YOLOv5

3、将CBAM集成到YOLOv5

4、用Bi-FPN替换PAN-Net

完整改进代码关注公众号后台回复:小目标检测

1 增加小目标检测层

YOLOv5小目标检测效果不好的一个原因是因为小目标样本的尺寸较小,而yolov5的下采样倍数比较大,较深的特征图很难学习到小目标的特征信息,因此提出增加小目标检测层对较浅特征图与深特征图拼接后进行检测。

这个方式的实现十分简单有效,只需要修改yolov5的模型文件yaml就可以增加小目标检测层,但是在增加检测层后,带来的问题就是计算量增加,导致推理检测速度降低。不过对于小目标,确实有很好的改善。

修改Anchor:增加一组较小的anchor

# anchors 
anchors:
  - [5,6, 8,14, 15,11]  #4
  - [10,13, 16,30, 33,23]  # P3/8
  - [30,61, 62,45, 59,119]  # P4/16
  - [116,90, 156,198, 373,326]  # P5/32

修改Head部分:

在第17层后,继续对特征图进行上采样等处理,使得特征图继续扩大,同时在第20层时,将获取到的大小为160X160的特征图与骨干网络中第2层特征图进行concat融合,以此获取更大的特征图进行小目标检测

# YOLOv5 head
   [-1, 1, Conv, [256, 1, 1]], #18  80*80
   [-1, 1, nn.Upsample, [None, 2, 'nearest']], #19  160*160
   [[-1, 2], 1, Concat, [1]], #20 cat backbone p2  160*160
   [-1, 3, BottleneckCSP, [256, False]], #21 160*160
# 篇幅限制中间省略一些代码,原代码关注公众后后台回复:YOLO小目标
   [[21, 24, 27, 30], 1, Detect, [nc, anchors]],  # Detect(p2, P3, P4, P5)
  ]

Transformer Prediction Heads (TPH)集成到YOLOv5

论文题目:TPH-YOLOv5: Improved YOLOv5 Based on Transformer Prediction Head for Object Detection on Drone-Captured Scenarios

TPH-YOLOv5的网络结构如下:

82c9bb782211cc6ef4d09186c35769db.png

TPH-YOLOv5用Transformer encoder块替换了YOLOv5中的一些卷积块和CSP bottleneck blocks,其结构如下图所示:每个Transformer encoder block包含2个子层。第1子层为multi-head attention layer,第2子层(MLP)为全连接层。每个子层之间使用残差连接。Transformer encoder block增加了捕获不同局部信息的能力。它还可以利用自注意力机制来挖掘特征表征潜能。

f11e6d593ba3cec602c61914ca381866.png

class C3TR(C3):
    # C3 module with TransformerBlock()
    def __init__(self, c1, c2, n=1, shortcut=True, g=1, e=0.5):
        super().__init__(c1, c2, n, shortcut, g, e)
        c_ = int(c2 * e)
        self.m = TransformerBlock(c_, c_, 4, n)
        
class TransformerBlock(nn.Module):
    # Vision Transformer https://arxiv.org/abs/2010.11929
    def __init__(self, c1, c2, num_heads, num_layers):
        super().__init__()
        self.conv = None
        if c1 != c2:
            self.conv = Conv(c1, c2)
        self.linear = nn.Linear(c2, c2)  # learnable position embedding
        self.tr = nn.Sequential(*(TransformerLayer(c2, num_heads) for _ in range(num_layers)))
        self.c2 = c2


    def forward(self, x):
        if self.conv is not None:
            x = self.conv(x)
        b, _, w, h = x.shape
        p = x.flatten(2).unsqueeze(0).transpose(0, 3).squeeze(3)
        return self.tr(p + self.linear(p)).unsqueeze(3).transpose(0, 3).reshape(b, self.c2, w, h)

与CSPDarknet53中的bottleneck blocks相比,Transformer encoder block可以捕获全局信息和丰富的上下文信息。

完整改进代码关注公众号后台回复:小目标检测

3、将CBAM集成到YOLOv5

在YOLOV5中插入CBAM,CBAM可以无缝地集成到任何CNN架构中,开销不会很大,而且可以与基本CNN网络一起进行端到端的训练。假设我们有一个中间特征F(H×W×C)作为卷积层的输出。根据层的深度,每个这样的特征层都会捕获有用的信息,例如简单的边缘、形状等,以获取更复杂的输入语义表示。我们希望网络更多地关注(即关注)这些特征图的重要部分。

15974eb9a11d3ab8cd617782bc355a77.png

class CBAM(nn.Module):
    def __init__(self, c1,c2):
        super(CBAM, self).__init__()
        self.channel_attention = ChannelAttentionModule(c1)
        self.spatial_attention = SpatialAttentionModule()


    def forward(self, x):
        out = self.channel_attention(x) * x
        out = self.spatial_attention(out) * out
        return out
        
class ChannelAttentionModule(nn.Module):
    def __init__(self, c1, reduction=16):
        super(ChannelAttentionModule, self).__init__()
        mid_channel = c1 // reduction
        self.avg_pool = nn.AdaptiveAvgPool2d(1)
        self.max_pool = nn.AdaptiveMaxPool2d(1)
        self.shared_MLP = nn.Sequential(
            nn.Linear(in_features=c1, out_features=mid_channel),
            nn.ReLU(),
            nn.Linear(in_features=mid_channel, out_features=c1)
        )
        self.sigmoid = nn.Sigmoid()
        #self.act=SiLU()
    def forward(self, x):
        avgout = self.shared_MLP(self.avg_pool(x).view(x.size(0),-1)).unsqueeze(2).unsqueeze(3)
        maxout = self.shared_MLP(self.max_pool(x).view(x.size(0),-1)).unsqueeze(2).unsqueeze(3)
        return self.sigmoid(avgout + maxout)
        
        
 class SpatialAttentionModule(nn.Module):
    def __init__(self):
        super(SpatialAttentionModule, self).__init__()
        self.conv2d = nn.Conv2d(in_channels=2, out_channels=1, kernel_size=7, stride=1, padding=3) 
        #self.act=SiLU()
        self.sigmoid = nn.Sigmoid()
    def forward(self, x):
        avgout = torch.mean(x, dim=1, keepdim=True)
        maxout, _ = torch.max(x, dim=1, keepdim=True)
        out = torch.cat([avgout, maxout], dim=1)
        out = self.sigmoid(self.conv2d(out))
        return out

使用CBAM可以提取注意区域,以帮助YOLOv5抵制令人困惑的信息,并关注有用的目标对象。

完整版改进代码关注公众号后台回复:小目标检测

4、用Bi-FPN替换PAN-Net

将位于Head和Backbone之间的结构称为“Neck”,其目标是在将Backbone提取的信息反馈到Head之前尽可能多地聚合这些信息。该结构通过防止小目标信息丢失,在传递小目标信息方面发挥了重要作用。它通过再次提高特征图的分辨率来做到这一点,这样来自Backbone的不同层的特征就可以被聚合,以提升整体的检测性能。

8733114f59da21283eefb523cd2ea441.png

Bi-FPN是在EfficientDet中提出了一种加权的双向特征金字塔网络,它允许简单和快速的多尺度特征融合。目的是为了追求更高效的多尺度融合方式。以往的特征融合是平等地对待不同尺度特征,Bi-FPN引入了权重(类似于attention),更好地平衡不同尺度的特征信息。

neck部分模型结构文件yaml代码

# YOLOv5 v6.0 head
head:
  [[-1, 1, Conv, [512, 1, 1]],
   [-1, 1, nn.Upsample, [None, 2, 'nearest']],
   [[-1, 6], 1, Concat, [1]],  # cat backbone P4
   [-1, 3, C3, [512, False]],  # 13


   [-1, 1, Conv, [256, 1, 1]],
   [-1, 1, nn.Upsample, [None, 2, 'nearest']],
   [[-1, 4], 1, Concat, [1]],  # cat backbone P3
   [-1, 3, C3, [256, False]],  # 17 (P3/8-small)
   
   [ -1, 1, Conv, [ 128, 1, 1 ] ],
   [ -1, 1, nn.Upsample, [ None, 2, 'nearest' ] ],
   [ [ -1, 2 ], 1, Concat, [ 1 ] ],  # cat backbone P2
   [-1, 1, SPP, [128, [5, 9, 13]]],
   [ -1, 3, C3, [ 128, False ] ],  #  (P2/4-xsmall)
   [-1, 1, CBAM, [128]],           # 23

完整版改进代码关注公众号后台回复:小目标检测

这篇文章介绍了四种基于YOLOV5小目标检测的网络结构优化方法,主要是通过Transformer encoder块替换了YOLOv5中的一些卷积块,以及增加小目标检测层,集成CBAM空间通道注意力模块和使用Bi-FPN替换PAN-FPN,还有一些针对小目标检测技巧的改进,后续有时间我们在做总结。

检测效果图:

8791df7c319c7de12963dc1f09b17703.png

参考资料:

[1] Zhu X, Lyu S, Wang X, et al. TPH-YOLOv5: Improved YOLOv5 Based on Transformer Prediction Head for Object Detection on Drone-captured Scenarios[C]//Proceedings of the IEEE/CVF International Conference on Computer Vision. 2021: 2778-2788.

[2] Benjumea A, Teeti I, Cuzzolin F, et al. YOLO-Z: Improving small object detection in YOLOv5 for autonomous vehicles[J]. arXiv preprint arXiv:2112.11798, 2021.

[3]  集智书童@公众号 详细解读TPH-YOLOv5 | 让目标检测任务中的小目标无处遁形

https://mp.weixin.qq.com/s/QPwvZSjHRQsnME8CjUg6UA

[4]  集智书童@公众号 记录修改YOLOv5以适应小目标检测的实验过程https://cloud.tencent.com/developer/article/1925516

[5]  Tz..@CSDN FPN以及其他结构FPN——Bi-FPN重点(仅供自己学习使用)https://blog.csdn.net/qq_41456654/article/details/115558583

[6]  小俊俊的博客 @CSDN yolov5增加小目标检测层

https://blog.csdn.net/weixin_41868104/article/details/111596851?spm=1001.2014.3001.5501

本文仅做学术分享,如有侵权,请联系删文。

3D视觉精品课程推荐:

1.面向自动驾驶领域的多传感器数据融合技术

2.面向自动驾驶领域的3D点云目标检测全栈学习路线!(单模态+多模态/数据+代码)
3.彻底搞透视觉三维重建:原理剖析、代码讲解、及优化改进
4.国内首个面向工业级实战的点云处理课程
5.激光-视觉-IMU-GPS融合SLAM算法梳理和代码讲解
6.彻底搞懂视觉-惯性SLAM:基于VINS-Fusion正式开课啦
7.彻底搞懂基于LOAM框架的3D激光SLAM: 源码剖析到算法优化
8.彻底剖析室内、室外激光SLAM关键算法原理、代码和实战(cartographer+LOAM +LIO-SAM)

9.从零搭建一套结构光3D重建系统[理论+源码+实践]

10.单目深度估计方法:算法梳理与代码实现

11.自动驾驶中的深度学习模型部署实战

12.相机模型与标定(单目+双目+鱼眼)

13.重磅!四旋翼飞行器:算法与实战

重磅!3DCVer-学术论文写作投稿 交流群已成立

扫码添加小助手微信,可申请加入3D视觉工坊-学术论文写作与投稿 微信交流群,旨在交流顶会、顶刊、SCI、EI等写作与投稿事宜。

同时也可申请加入我们的细分方向交流群,目前主要有3D视觉CV&深度学习SLAM三维重建点云后处理自动驾驶、多传感器融合、CV入门、三维测量、VR/AR、3D人脸识别、医疗影像、缺陷检测、行人重识别、目标跟踪、视觉产品落地、视觉竞赛、车牌识别、硬件选型、学术交流、求职交流、ORB-SLAM系列源码交流、深度估计等微信群。

一定要备注:研究方向+学校/公司+昵称,例如:”3D视觉 + 上海交大 + 静静“。请按照格式备注,可快速被通过且邀请进群。原创投稿也请联系。

8e34ae1360321dc1ece30cd85f2c2313.png

▲长按加微信群或投稿

ae908dc90a775d56f4c98da93ee5090b.png

▲长按关注公众号

3D视觉从入门到精通知识星球:针对3D视觉领域的视频课程(三维重建系列三维点云系列结构光系列手眼标定相机标定激光/视觉SLAM自动驾驶等)、知识点汇总、入门进阶学习路线、最新paper分享、疑问解答五个方面进行深耕,更有各类大厂的算法工程人员进行技术指导。与此同时,星球将联合知名企业发布3D视觉相关算法开发岗位以及项目对接信息,打造成集技术与就业为一体的铁杆粉丝聚集区,近4000星球成员为创造更好的AI世界共同进步,知识星球入口:

学习3D视觉核心技术,扫描查看介绍,3天内无条件退款

9764c49d70a520ec080d8a48a7281e9d.png

 圈里有高质量教程资料、答疑解惑、助你高效解决问题

觉得有用,麻烦给个赞和在看~  

评论 36
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值