基于改进YOLOv5的小目标检测 | 添加CBAM注意机制 + 更换Neck网络之BiFPN + 增加高分辨率检测头

前言：Hello大家好，我是小哥谈。本文针对图像中小目标难以检测的问题，提出了一种基于YOLOv5的改进模型。在主干网络中，加入CBAM注意力模块增强网络特征提取能力；在颈部网络部分，使用BiFPN结构替换PANet结构，强化底层特征利用；在检测头部分，增加高分辨率检测头，改善对于微小目标的检测能力。本文算法在无人机数据集VisDrone2019数据集上进行了多次对比实验，结果表明本文算法可以有效地检测小目标。🌈

🚀1.基础概念

🚀2.网络结构

🚀1.基础概念

目标检测是计算机视觉的重要研究方向，也是众多复杂视觉任务的基础，被广泛应用于工业、农业等领域。尽管深度学习技术的出现使得目标检测取得了较大的突破，但是现有的方法依然很难较好的检测小目标。小目标尺寸小、特征不明显，在检测中误检率、漏检率一般均较高，因此提升小目标检测性能仍然是一个具有挑战性的研究方向。

本节贡献：

针对小目标对象提出了一种新的检测模型。
基于注意力机制，通过改进主干网络增强模型对于目标特征的提取。
基于双向交叉连接和加权融合的思想，引入额外特征分支，强化底层特征利用。
针对目标尺寸较小的特点，在检测头部分增加高分辨率检测分支，增强模型对于微小目标的检测能力。

小目标检测头：

原始YOLOv5模型主干网络一共进行5次下采样，得到5层特征表达（P1、P2、P3、P4和P5）。尽管在颈部网络中，通过自上而下和自下而上的聚合路径实现了多尺度特征融合，但是并不影响特征图的尺度，最后检测头部分在通过P3、P4和P5这3级特征图引出的检测头上进行目标的检测，其特征图尺度分别为80×80、40×40、 20×20。为方便表达，通过Pi层特征图引出的检测头，以下简称为Pi层检测头。在小目标检测任务中，往往存在非常小的待检目标。在VisDrone2019公开无人机数据集中，甚至含有较多小于3×3像素的目标。这样的目标在经过多次下采样之后，其大部分特征信息已经丢失，尽管通过具有较高分辨率的P3层检测头依然难以检测到。

为了实现上述微小目标同样可以达到较好的检测效果，我们在YOLOv5模型上通过P2层特征引出了新的检测头。P2层检测头分辨率为160×160像素，相当于在主干网络中只进行了2次下采样操作，含有目标更为丰富的底层特征信息。颈部网络中自上而下和自下而上得到的两个P2层特征与主干网络中的同尺度特征通过concat形式进行特征融合，输出的特征为3个输入特征的融合结果，这样使得P2层检测头应对微小目标时，能够快速有效的检测。P2层检测头加上原始的3个检测头，可以有效缓解尺度方差所带来的负面影响，增加的检测头是针对底层特征的，是通过低水平、高分辨率的特征图生成的，该检测头对微小目标更加敏感。尽管添加这个检测头增加了模型的计算量和内存开销，但是对于微小目标的检测能力有着不小的提升。

CBAM注意力模块：

CBAM是一种轻量的注意力模块，其简单有效，可以直接集成到CNN架构中，并且可以端到端的对其进行训练。在给定特征映射的情况下，CBAM会依次沿着通道和空间两个独立维度推导注意映射，然后将注意映射与输入特征映射相乘，进行自适应特征细化。CBAM 模块的结构如下图所示，CBAM模块被集成到不同数据集和不同分类任务的不同模型中，模型性能均得到了较大提升，证明了CBAM模块的有效性。

小目标尺寸较小、特征少且不明显，在主干网络中加入CBAM注意力模块，可以增强网络对于目标特征的提取能力，同时直接改善颈部网络部分的特征融合。在检测任务中，CBAM注意力模块可以帮助模型有效的提取注意区域，提高检测性能。

BiFPN网络：

BiFPN网络是集成双向交叉连接和加权融合的一种高效的多尺度特征融合方法。自FPN被提出以来，FPN被广泛应用于多尺度特征融合。近些年来，PANet、M2det和NAS-FPN等更多的多尺度特征融合网络结构被研究学者们提出来，但是在融合不同层次的输入特征时，大部分工作都是不加区分的总结它们。然而，这些不同的输入特征有着不同分辨率，对融合的输出特征具有在不同的贡献。为此，提出了一种简单但高效的加权双向特征金字塔网络（BiFPN），它引入了可学习的权重因子来表征不同输入特征的重要程度，同时反复应用自顶向下和自底向上的多尺度特征融合。

为充分利用目标的底层特征，本节课对颈部网络进行了改进，将原始PANet网络换成BiFPN网络，以提高检测精度，其结构如下图所示。尽管YOLOv5中的PANet通过自顶向下和自底向上的路径聚合实现了较好的多尺度特征融合结果，但计算量较大，且自底向上特征融合阶段的输入特征中并没有主干网络生成的原始输出特征。BiFPN采用跨连接去除PANet中对特征融合贡献度较小的节点，在同一尺度的输入节点和输出节点间增加一个跳跃连接，在不增加较多成本的同时，融合了更多的特征。在同一特征尺度上，把每一个双向路径看作一个特征网络层，并多次反复利用同一层，以实现更高层次的特征融合。

通过上述三种改进方式后，本节课针对YOLOv5改进后的网络结构图如下所示：

🚀2.网络结构

本文的改进是基于YOLOv5-6.0版本，关于其网络结构具体如下图所示：

本文对YOLOv5的改进是添加CBAM注意机制 + 更换Neck网络之BiFPN + 增加高分辨率检测头，改进后的网络结构图具体如下图所示：

🚀3.添加步骤

针对本文的改进，具体步骤如下所示：👇

步骤1：common.py文件修改

步骤2：yolo.py文件修改

步骤3：创建自定义yaml文件

步骤4：修改自定义yaml文件

步骤5：验证是否加入成功

步骤6：修改默认参数

步骤7：实际训练测试

🚀4.改进方法

🍀🍀步骤1：common.py文件修改

在common.py中添加CBAM注意力机制和BiFPN模块代码，所要添加模块的代码如下所示，将其复制粘贴到common.py文件末尾的位置。

# 基于改进YOLOv5的小目标检测
# By CSDN 小哥谈
# CBAM注意力机制代码
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.LeakyReLU(0.1, inplace=True),
            nn.Linear(in_features=mid_channel, out_features=c1)
        )
        self.act = nn.Sigmoid()
        # self.act=nn.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.act(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 = 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.act(self.conv2d(out))
        return out

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

# BiFPN模块代码
class BiFPN_Add2(nn.Module):
    def __init__(self, c1, c2):
        super(BiFPN_Add2, self).__init__()
        # 设置可学习参数 nn.Parameter的作用是：将一个不可训练的类型Tensor转换成可以训练的类型parameter
        # 并且会向宿主模型注册该参数 成为其一部分 即model.parameters()会包含这个parameter
        # 从而在参数优化的时候可以自动一起优化
        self.w = nn.Parameter(torch.ones(2, dtype=torch.float32), requires_grad=True)
        self.epsilon = 0.0001
        self.conv = nn.Conv2d(c1, c2, kernel_size=1, stride=1, padding=0)
        self.silu = nn.SiLU()

    def forward(self, x):
        w = self.w
        weight = w / (torch.sum(w, dim=0) + self.epsilon)
        return self.conv(self.silu(weight[0] * x[0] + weight[1] * x[1]))

# 三个特征图add操作
class BiFPN_Add3(nn.Module):
    def __init__(self, c1, c2):
        super(BiFPN_Add3, self).__init__()
        self.w = nn.Parameter(torch.ones(3, dtype=torch.float32), requires_grad=True)
        self.epsilon = 0.0001
        self.conv = nn.Conv2d(c1, c2, kernel_size=1, stride=1, padding=0)
        self.silu = nn.SiLU()

    def forward(self, x):
        w = self.w
        weight = w / (torch.sum(w, dim=0) + self.epsilon)
        # Fast normalized fusion
        return self.conv(self.silu(weight[0] * x[0] + weight[1] * x[1] + weight[2] * x[2]))

🍀🍀步骤2：yolo.py文件修改

在yolo.py文件中找到parse_model函数，在下图中所示位置添加CBAM，并且添加下列代码：

# -----------BiFPN-------------------
        elif m in [BiFPN_Add2, BiFPN_Add3]:
            c2 = max([ch[x] for x in f])
        # ------------end-----------------------

具体添加位置如下图所示：

🍀🍀步骤3：创建自定义yaml文件

在models文件夹中复制yolov5s.yaml，粘贴并重命名为：yolov5s_CBAM_BiFPN.yaml。具体如下图所示：

🍀🍀步骤4：修改自定义yaml文件

本步骤是修改yolov5s_CBAM_BiFPN.yaml，根据改进后的网络结构图进行修改。

修改后的完整yaml文件如下所示：

# YOLOv5 🚀 by Ultralytics, GPL-3.0 license
# By CSDN 小哥谈

# Parameters
nc: 80 # number of classes
depth_multiple: 0.33 # model depth multiple
width_multiple: 0.50 # layer channel multiple
anchors:
  - [5, 6, 8, 14, 15, 11]  # P2/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

backbone:
  # [from, number, module, args]
  [
    [-1, 1, Conv, [64, 6, 2, 2]], # 0-P1/2
    [-1, 1, Conv, [128, 3, 2]], # 1-P2/4
    [-1, 3, C3, [128]],  # 2 160 * 160
    [-1, 1, CBAM, [128]],  # 3
    [-1, 1, Conv, [256, 3, 2]], # 4-P3/8
    [-1, 6, C3, [256]],  # 80 * 80
    [-1, 1, Conv, [512, 3, 2]], # 6-P4/16
    [-1, 9, C3, [512]],  # 7 40 * 40
    [-1, 1, Conv, [1024, 3, 2]], # 8-P5/32
    [-1, 3, C3, [1024]],  # 9 20 * 20
    [-1, 1, SPPF, [1024, 5]], # 10
  ]
head: [
    [-1, 1, Conv, [512, 1, 1]], # 11
    [-1, 1, nn.Upsample, [None, 2, "nearest"]], # 12
    [[-1, 7], 1, BiFPN_Add2, [256, 256]], # cat backbone P4  # 13
    [-1, 3, C3, [512, False]], # 14

    [-1, 1, Conv, [256, 1, 1]], # 15
    [-1, 1, nn.Upsample, [None, 2, "nearest"]], # 16
    [[-1, 5], 1, BiFPN_Add2, [128, 128]], # cat backbone P3 # 17
    [-1, 3, C3, [256, False]], # (P3/8-small) # 18

    [-1, 1, Conv, [128, 1, 1]], # 19
    [-1, 1, nn.Upsample, [None, 2, "nearest"]], #20
    [[-1, 2], 1, BiFPN_Add2, [64, 64]],  # 21  p2
    [-1, 3, C3, [128, False]],  # 22

    [-1, 1, Conv, [256, 3, 2]],  # 23
    [[-1, 18, 5], 1, BiFPN_Add3, [128, 128]], # 24 p3
    [-1, 3, C3, [256, False]], # 25 p3

    [-1, 1, Conv, [512, 3, 2]], # 26
    [[-1, 14, 7], 1, BiFPN_Add3, [256, 256]], # 27 cat head P4
    [-1, 3, C3, [512, False]], # 28 (P4/16-medium)

    [-1, 1, Conv, [512, 3, 2]],  # 29
    [[-1, 11], 1, BiFPN_Add2, [256, 256]], # 30 cat head P5
    [-1, 3, C3, [1024, False]], # 31 (P5/32-large)

    [[22, 25, 28, 31], 1, Detect, [nc, anchors]], # Detect(P2, P3, P4, P5)
  ]

🍀🍀步骤5：验证是否加入成功

在yolo.py文件里，将配置改为我们刚才自定义的yolov5s_CBAM_BiFPN.yaml。

修改1，位置位于yolo.py文件165行左右，具体如图所示：

修改2，位置位于yolo.py文件363行左右，具体如下图所示：

配置完毕之后，点击“运行”，结果如下图所示：

由运行结果可知，与我们前面更改后的网络结构图相一致，证明添加成功了！✅

参数量对比：

yolov5s.yaml：214 layers, 7235389 parameters, 7235389 gradients, 16.6 GFLOPs

yolov5s_CBAM_BiFPN.yaml：284 layers, 7597553 parameters, 7597553 gradients, 21.5 GFLOPs

🍀🍀步骤6：修改默认参数

修改1，设置可学习权重，将下列代码放置在train.py文件中（160行左右）。

# 设置可学习权重
    g0, g1, g2 = [], [], []  # optimizer parameter groups
    for v in model.modules():
        # hasattr: 测试指定的对象是否具有给定的属性，返回一个布尔值
        if hasattr(v, 'bias') and isinstance(v.bias, nn.Parameter):  # bias
            g2.append(v.bias)  # biases
        if isinstance(v, nn.BatchNorm2d):  # weight (no decay)
            g0.append(v.weight)
        elif hasattr(v, 'weight') and isinstance(v.weight, nn.Parameter):
            g1.append(v.weight)
        # BiFPN_Concat
        elif isinstance(v, BiFPN_Add2) and hasattr(v, 'w') and isinstance(v.w, nn.Parameter):
            g1.append(v.w)
        elif isinstance(v, BiFPN_Add3) and hasattr(v, 'w') and isinstance(v.w, nn.Parameter):
            g1.append(v.w)

具体放置位置如下图所示：

修改2，在train.py文件中找到parse_opt函数，然后将第二行 '--cfg' 的default改为 ' models/yolov5s_CBAM_BiFPN.yaml '，然后就可以开始进行训练了。🎈🎈🎈

🍀🍀步骤7：实际训练测试

在本步骤中，parse_opt函数中的参数'--weights'采用的是yolov5s.pt，'--data'所采用的是helmet.yaml（作者提前创建的安全帽佩戴检测地址及分类信息，同学可自定义），然后设置'--epochs'为100轮。相关参数设置完毕后，点击运行train.py文件，没有发生报错，模型正常训练，具体如下图所示：👇

说明：后期实际训练时，根据论文所设置的指标进行训练，上述只是进行测试。

🚀5.试验分析

🍀🍀5.1 数据集

VisDrone2019数据集由天津大学机器学习和数据挖掘实验室AISKYEYE团队收集，基准数据集包含288个视频片段、261908帧和10209幅静态图像。数据集是在不同的场景、不同的天气和光照条件下使用各种无人机平台（即不同型号的无人机）收集的，手工标注了超过260万个边界框或经常感兴趣的目标点，包括行人、人、轿车、货车、公共汽车、卡车、摩托车、自行车、遮阳篷三轮车和三轮车等10个感兴趣的对象类别。其各类标签数量如表1所示。