当前位置：首页 > news >正文

YOLOv5改进CBAM【保姆级教程】

news 来源：原创 2025/4/22 8:32:14

1. 前言🐱‍🏍

YOLOv5是目前目标检测任务中非常流行的一种高效模型，广泛应用于实时目标检测中。尽管其在速度和精度方面都有着优异的表现，但在复杂背景和小目标检测方面，YOLOv5仍然存在一定的提升空间。为了解决这一问题，将CBAM（Convolutional Block Attention Module）注意力机制引入YOLOv5模型的方法，从而提升其对重要特征的关注能力，提高检测精度。

2. CBAM注意力机制🚀🚀

CBAM（Convolutional Block Attention Module）是一种轻量级的注意力机制，旨在帮助神经网络学习到更有用的特征信息。它由两个主要模块组成：通道注意力模块和空间注意力模块。

2.1 通道注意力机制

通道注意力机制的目的是通过加权各个通道的特征图，帮助模型集中精力关注重要通道。其计算过程如下：

（1）全局池化操作：对输入特征图 𝑋进行全局平均池化和全局最大池化操作，得到两个池化结果：
$F_{\text{avg}} = \text{AvgPool}(X), \quad F_{\text{max}} = \text{MaxPool}(X)$
其中，𝐹avg 和𝐹max是通过全局平均池化和全局最大池化得到的特征图。

（2）通道加权：通过一个共享的全连接网络（MLP）对这两个池化结果进行融合，并得到每个通道的权重：
$M_c = \sigma\left(W_2 \delta\left(W_1 \left[ F_{\text{avg}}, F_{\text{max}} \right]\right)\right)$
其中，𝑊1和 𝑊2是可训练的权重矩阵，𝛿 是ReLU激活函数，𝜎是sigmoid激活函数，𝑀𝑐是通道注意力权重。

2.2 空间注意力机制

空间注意力机制的目的是通过加权输入特征图的空间区域，使模型能够关注到关键的空间位置。空间注意力的计算公式如下：
$M_s = \sigma\left(f_{7 \times 7} \left([ F_{\text{avg}}, F_{\text{max}} ]\right)\right)$

其中，𝑓7×7是一个7x7的卷积操作，[𝐹avg,𝐹max]是通道平均池化和最大池化结果的拼接，𝑀𝑠是空间注意力权重。

2.3 CBAM模块的结合📈

CBAM模块通过结合通道和空间两个方面的注意力机制，能够更加精细地对特征图进行加权，从而提升模型的性能。通过对输入特征图 𝑋进行加权，CBAM模块的输出 𝑋out 计算公式为：
$X_{\text{out}} = X \cdot M_c \cdot M_s$
其中，𝑀𝑐是通道注意力权重，𝑀𝑠是空间注意力权重，⋅ 表示逐元素相乘操作。

3. 💻代码实现💻

3.1 修改common.py文件

（1）找到代码中的common.py文件，如下图所示：

在这里插入图片描述

（2）将以下代码复制粘贴到common.py中最底部。

###CBAM注意力机制# 通道注意力模块
class ChannelAttention(nn.Module):def __init__(self, in_planes, ratio=16):super(ChannelAttention, self).__init__()self.avg_pool = nn.AdaptiveAvgPool2d(1)  # 自适应平均池化self.max_pool = nn.AdaptiveMaxPool2d(1)  # 自适应最大池化# 两个卷积层用于从池化后的特征中学习注意力权重self.fc1 = nn.Conv2d(in_planes, in_planes // ratio, 1, bias=False)  # 第一个卷积层，降维self.relu1 = nn.ReLU()  # ReLU激活函数self.fc2 = nn.Conv2d(in_planes // ratio, in_planes, 1, bias=False)  # 第二个卷积层，升维self.sigmoid = nn.Sigmoid()  # Sigmoid函数生成最终的注意力权重def forward(self, x):avg_out = self.fc2(self.relu1(self.fc1(self.avg_pool(x))))  # 对平均池化的特征进行处理max_out = self.fc2(self.relu1(self.fc1(self.max_pool(x))))  # 对最大池化的特征进行处理out = avg_out + max_out  # 将两种池化的特征加权和作为输出return self.sigmoid(out)  # 使用sigmoid激活函数计算注意力权重# 空间注意力模块
class SpatialAttention(nn.Module):def __init__(self, kernel_size=7):super(SpatialAttention, self).__init__()assert kernel_size in (3, 7), 'kernel size must be 3 or 7'  # 核心大小只能是3或7padding = 3 if kernel_size == 7 else 1  # 根据核心大小设置填充# 卷积层用于从连接的平均池化和最大池化特征图中学习空间注意力权重self.conv1 = nn.Conv2d(2, 1, kernel_size, padding=padding, bias=False)self.sigmoid = nn.Sigmoid()  # Sigmoid函数生成最终的注意力权重def forward(self, x):avg_out = torch.mean(x, dim=1, keepdim=True)  # 对输入特征图执行平均池化max_out, _ = torch.max(x, dim=1, keepdim=True)  # 对输入特征图执行最大池化x = torch.cat([avg_out, max_out], dim=1)  # 将两种池化的特征图连接起来x = self.conv1(x)  # 通过卷积层处理连接后的特征图return self.sigmoid(x)  # 使用sigmoid激活函数计算注意力权重class CBAM(nn.Module):def __init__(self, in_planes, ratio=16, kernel_size=7):super(CBAM, self).__init__()self.ca = ChannelAttention(in_planes, ratio)  # 通道注意力实例self.sa = SpatialAttention(kernel_size)  # 空间注意力实例def forward(self, x):out = x * self.ca(x)  # 使用通道注意力加权输入特征图result = out * self.sa(out)  # 使用空间注意力进一步加权特征图return result  # 返回最终的特征图

（3）插入后如图所示：

在这里插入图片描述

3.2 修改yolo.py文件

（1）找到代码中的yolo.py文件，如下图所示：

在这里插入图片描述

（2）在yolo.py中找到from models.common import加入CBAM模块：

from models.common import (C3,C3SPP,C3TR,SPP,SPPF,Bottleneck,BottleneckCSP,C3Ghost,C3x,Classify,Concat,Contract,Conv,CrossConv,DetectMultiBackend,DWConv,DWConvTranspose2d,Expand,Focus,GhostBottleneck,GhostConv,Proto,CBAM,
)

（3）找到**def parse_model(d, ch)😗*函数，添加CBAM模块，直接复制以下代码，替换源代码中的此函数：

def parse_model(d, ch):"""Parses a YOLOv5 model from a dict `d`, configuring layers based on input channels `ch` and model architecture."""LOGGER.info(f"\n{'':>3}{'from':>18}{'n':>3}{'params':>10}  {'module':<40}{'arguments':<30}")anchors, nc, gd, gw, act, ch_mul = (d["anchors"],d["nc"],d["depth_multiple"],d["width_multiple"],d.get("activation"),d.get("channel_multiple"),)if act:Conv.default_act = eval(act)  # redefine default activation, i.e. Conv.default_act = nn.SiLU()LOGGER.info(f"{colorstr('activation:')} {act}")  # printif not ch_mul:ch_mul = 8na = (len(anchors[0]) // 2) if isinstance(anchors, list) else anchors  # number of anchorsno = na * (nc + 5)  # number of outputs = anchors * (classes + 5)layers, save, c2 = [], [], ch[-1]  # layers, savelist, ch outfor i, (f, n, m, args) in enumerate(d["backbone"] + d["head"]):  # from, number, module, argsm = eval(m) if isinstance(m, str) else m  # eval stringsfor j, a in enumerate(args):with contextlib.suppress(NameError):args[j] = eval(a) if isinstance(a, str) else a  # eval stringsn = n_ = max(round(n * gd), 1) if n > 1 else n  # depth gainif m in {Conv,GhostConv,Bottleneck,GhostBottleneck,SPP,SPPF,DWConv,MixConv2d,Focus,CrossConv,BottleneckCSP,C3,C3TR,C3SPP,C3Ghost,nn.ConvTranspose2d,DWConvTranspose2d,C3x,CBAM,}:c1, c2 = ch[f], args[0]if c2 != no:  # if not outputc2 = make_divisible(c2 * gw, ch_mul)args = [c1, c2, *args[1:]]if m in {BottleneckCSP, C3, C3TR, C3Ghost, C3x}:args.insert(2, n)  # number of repeatsn = 1elif m is nn.BatchNorm2d:args = [ch[f]]elif m is Concat:c2 = sum(ch[x] for x in f)# TODO: channel, gw, gdelif m in {Detect, Segment}:args.append([ch[x] for x in f])if isinstance(args[1], int):  # number of anchorsargs[1] = [list(range(args[1] * 2))] * len(f)if m is Segment:args[3] = make_divisible(args[3] * gw, ch_mul)elif m is Contract:c2 = ch[f] * args[0] ** 2elif m is Expand:c2 = ch[f] // args[0] ** 2else:c2 = ch[f]m_ = nn.Sequential(*(m(*args) for _ in range(n))) if n > 1 else m(*args)  # modulet = str(m)[8:-2].replace("__main__.", "")  # module typenp = sum(x.numel() for x in m_.parameters())  # number paramsm_.i, m_.f, m_.type, m_.np = i, f, t, np  # attach index, 'from' index, type, number paramsLOGGER.info(f"{i:>3}{str(f):>18}{n_:>3}{np:10.0f}  {t:<40}{str(args):<30}")  # printsave.extend(x % i for x in ([f] if isinstance(f, int) else f) if x != -1)  # append to savelistlayers.append(m_)if i == 0:ch = []ch.append(c2)return nn.Sequential(*layers), sorted(save)

3.2 修改yaml文件

（1）新建一个yolov5_CBAM.yaml文件，复制以下代码粘贴进去：

# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license# Parameters
nc: 80 # number of classes
depth_multiple: 0.33 # model depth multiple
width_multiple: 0.50 # layer channel multiple
anchors:- [10, 13, 16, 30, 33, 23] # P3/8- [30, 61, 62, 45, 59, 119] # P4/16- [116, 90, 156, 198, 373, 326] # P5/32# YOLOv5 v6.0 backbone
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]],[-1, 1, Conv, [256, 3, 2]], # 3-P3/8[-1, 6, C3, [256]],[-1, 1, Conv, [512, 3, 2]], # 5-P4/16[-1, 9, C3, [512]],[-1, 1, Conv, [1024, 3, 2]], # 7-P5/32[-1, 3, C3, [1024]],[-1, 1, CBAM, [1024]],       # 9-CBAM  [-1, 1, SPPF, [1024, 5]], # 9]# 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, [256, 3, 2]],[[-1, 14], 1, Concat, [1]], # cat head P4[-1, 3, C3, [512, False]], # 20 (P4/16-medium)[-1, 1, Conv, [512, 3, 2]],[[-1, 10], 1, Concat, [1]], # cat head P5[-1, 3, C3, [1024, False]], # 23 (P5/32-large)[[17, 20, 23], 1, Detect, [nc, anchors]], # Detect(P3, P4, P5)]