当前位置:   article > 正文

YOLOv5改进 | 注意力机制 | 通道和空间的双重作用的CBAM注意力机制_通道注意力加空间注意力模块

通道注意力加空间注意力模块

深度学习目标检测领域,YOLOv5成为了备受关注的模型之一。本文给大家带来的是通道和空间的双重作用的CBAM注意力机制。文章在介绍主要的原理后,将手把手教学如何进行模块的代码添加和修改,并将修改后的完整代码放在文章的最后,方便大家一键运行,小白也可轻松上手实践。以帮助您更好地学习深度学习目标检测YOLO系列的挑战。


专栏地址YOLOv5改进+入门——持续更新各种有效涨点方法 

目录

1.原理 

2.YOLOv5添加CBAM注意力机制

2.1 CBAM注意力机制代码

2.2新增yaml文件

2.3 注册模块

2.4 执行程序

3.总结 


1.原理 

论文地址:CBAM: Convolutional Block Attention Module点击即可跳转

实现代码:CBAM代码实现点击即可跳转

CBAM(Convolutional Block Attention Module)是一种引入了注意力机制的卷积神经网络模块,旨在增强CNN模型的表征能力和性能。它由两个关键组件组成:通道注意力模块(Channel Attention Module)和空间注意力模块(Spatial Attention Module)。

1. 通道注意力模块(CAM):
   CAM主要用于建模特征图在通道维度上的关系。它的目标是学习每个通道的重要性,并对不同通道的特征进行加权,以提升有用特征的影响力,抑制无用特征的干扰。
   CAM首先对输入的特征图进行全局平均池化(Global Average Pooling),将每个通道的特征图压缩成一个标量,然后通过全连接层(FC)学习得到每个通道的权重。这些权重用于对每个通道的特征图进行加权,得到加权后的特征表示。

2. 空间注意力模块(SAM):
   SAM用于捕捉特征图在空间维度上的重要性。它的目标是学习不同空间位置的权重,使网络能够更好地关注图像中的重要区域。
   SAM首先对特征图进行两种池化操作:最大池化和平均池化。这两种池化操作分别用于捕捉特征图中的局部显著性和全局分布信息。然后,将两种池化结果结合,并通过全连接层学习得到每个空间位置的权重,以产生最终的空间注意力图。
   
通过结合通道注意力和空间注意力,CBAM可以使网络更好地理解输入数据中的关键信息,并提高模型在各种视觉任务上的性能。这种注意力机制的引入使得网络能够自适应地调整特征图中不同通道和空间位置的重要性,从而有效地提升了模型的表现力和泛化能力。CBAM已经被成功应用于图像分类、目标检测、语义分割等多个计算机视觉任务中,取得了显著的性能提升。

CBAM结构简图

2.YOLOv5添加CBAM注意力机制

2.1 CBAM注意力机制代码

关键步骤一:将下面代码添加到 yolov5/models/common.py中任意位置

  1. class ChannelAttention(nn.Module):
  2. def __init__(self, in_planes, ratio=16):
  3. super(ChannelAttention, self).__init__()
  4. self.avg_pool = nn.AdaptiveAvgPool2d(1)
  5. self.max_pool = nn.AdaptiveMaxPool2d(1)
  6. self.f1 = nn.Conv2d(in_planes, in_planes // ratio, 1, bias=False)
  7. self.relu = nn.ReLU()
  8. self.f2 = nn.Conv2d(in_planes // ratio, in_planes, 1, bias=False)
  9. self.sigmoid = nn.Sigmoid()
  10. def forward(self, x):
  11. avg_out = self.f2(self.relu(self.f1(self.avg_pool(x))))
  12. max_out = self.f2(self.relu(self.f1(self.max_pool(x))))
  13. out = self.sigmoid(avg_out + max_out)
  14. return out
  15. class SpatialAttention(nn.Module):
  16. def __init__(self, kernel_size=7):
  17. super(SpatialAttention, self).__init__()
  18. assert kernel_size in (3, 7), 'kernel size must be 3 or 7'
  19. padding = 3 if kernel_size == 7 else 1
  20. # (特征图的大小-算子的size+2*padding)/步长+1
  21. self.conv = nn.Conv2d(2, 1, kernel_size, padding=padding, bias=False)
  22. self.sigmoid = nn.Sigmoid()
  23. def forward(self, x):
  24. # 1*h*w
  25. avg_out = torch.mean(x, dim=1, keepdim=True)
  26. max_out, _ = torch.max(x, dim=1, keepdim=True)
  27. x = torch.cat([avg_out, max_out], dim=1)
  28. #2*h*w
  29. x = self.conv(x)
  30. #1*h*w
  31. return self.sigmoid(x)
  32. class CBAM(nn.Module):
  33. # CSP Bottleneck with 3 convolutions
  34. def __init__(self, c1, c2, ratio=16, kernel_size=7): # ch_in, ch_out, number, shortcut, groups, expansion
  35. super(CBAM, self).__init__()
  36. self.channel_attention = ChannelAttention(c1, ratio)
  37. self.spatial_attention = SpatialAttention(kernel_size)
  38. def forward(self, x):
  39. out = self.channel_attention(x) * x
  40. # c*h*w
  41. # c*h*w * 1*h*w
  42. out = self.spatial_attention(out) * out
  43. return out

 CBAM(Convolutional Block Attention Module)注意力机制的流程可以总结如下:

1. 输入特征图:接收来自上一层或输入图像的特征图作为输入。

2. 通道注意力模块(Channel Attention Module,CAM):
   对输入特征图进行全局平均池化(Global Average Pooling),将每个通道的特征进行降维,得到每个通道的全局描述。
   通过全连接层(Fully Connected Layer)学习得到每个通道的权重向量,这些权重用于衡量每个通道的重要性。
   将学习到的权重与原始特征图相乘,以加权增强有用特征和抑制无用特征。

3. 空间注意力模块(Spatial Attention Module,SAM):
   对输入特征图进行最大池化(Max Pooling)和平均池化(Average Pooling),分别捕获局部显著性和全局分布信息。
   将两种池化结果进行组合(如相加),得到综合的空间注意力图。
   通过激活函数(如sigmoid)对空间注意力图进行归一化,得到每个空间位置的权重。

4. 结合通道和空间注意力:
   将通道注意力加权后的特征图与空间注意力加权后的特征图进行逐元素相乘,得到最终的注意力增强特征图。

5. 输出:最终的注意力增强特征图作为模块的输出,传递给下一层网络进行后续的处理,如分类、检测或分割等任务。

整个CBAM注意力机制的流程是将通道注意力和空间注意力相结合,使得网络能够自适应地调整不同通道和空间位置的重要性,从而提升模型的性能和泛化能力。

2.2新增yaml文件

关键步骤二:在 /yolov5/models/ 下新建文件 yolov5_cbam.yaml并将下面代码复制进去

  1. # YOLOv5
    声明:本文内容由网友自发贡献,不代表【wpsshop博客】立场,版权归原作者所有,本站不承担相应法律责任。如您发现有侵权的内容,请联系我们。转载请注明出处:https://www.wpsshop.cn/w/Guff_9hys/article/detail/873844
    推荐阅读
    相关标签