一、CBAM注意力机制
CBAM(Convolutional Block Attention Module)是一种用于增强卷积神经网络(CNN)性能的注意力机制模块。它由Sanghyun Woo等人在2018年的论文(https://arxiv.org/abs/1807.06521)中提出。CBAM的主要目标是通过在CNN中引入通道注意力和空间注意力来提高模型的感知能力,从而在不增加网络复杂性的情况下改善性能。
CBAM由两个关键部分组成:通道注意力模块(C-channel)和空间注意力模块(S-channel)。这两个模块可以分别嵌入到CNN中的不同层,以增强特征表示。
二、通道注意力模块
通道注意力模块的目标是增强每个通道的特征表达
三、空间注意力模块
空间注意力模块的目标是强调图像中不同位置的重要性
四、CBAM模块
CBAM就是将通道注意力模块和空间注意力模块的输出特征逐元素相乘,得到最终的注意力增强特征。这个增强的特征将用作后续网络层的输入,以在保留关键信息的同时,抑制噪声和无关信息。原文实验证明先进行通道维度的整合,再进行空间维度的整合,模型效果更好
五、如何在YOLOv8中添加CBAM注意力模块
CBAM模块代码(pytorch):
#CBAM注意力
class ChannelAttention(nn.Module):
"""Channel-attention module https://github.com/open-mmlab/mmdetection/tree/v3.0.0rc1/configs/rtmdet."""
def __init__(self, channels: int) -> None:
super().__init__()
self.pool = nn.AdaptiveAvgPool2d(1)
self.fc = nn.Conv2d(channels, channels, 1, 1, 0, bias=True)
self.act = nn.Sigmoid()
def forward(self, x: torch.Tensor) -> torch.Tensor:
return x * self.act(self.fc(self.pool(x)))
class SpatialAttention(nn.Module):
"""Spatial-attention module."""
def __init__(self, kernel_size=7):
"""Initialize Spatial-attention module with kernel size argument."""
super().__init__()
assert kernel_size in (3, 7), 'kernel size must be 3 or 7'
padding = 3 if kernel_size == 7 else 1
self.cv1 = nn.Conv2d(2, 1, kernel_size, padding=padding, bias=False)
self.act = nn.Sigmoid()
def forward(self, x):
"""Apply channel and spatial attention on input for feature recalibration."""
return x * self.act(self.cv1(torch.cat([torch.mean(x, 1, keepdim=True), torch.max(x, 1, keepdim=True)[0]], 1)))
class CBAM(nn.Module):
"""Convolutional Block Attention Module."""
def __init__(self, c1, kernel_size=7): # ch_in, kernels
super().__init__()
self.channel_attention = ChannelAttention(c1)
self.spatial_attention = SpatialAttention(kernel_size)
def forward(self, x):
"""Applies the forward pass through C1 module."""
return self.spatial_attention(self.channel_attention(x))
首先将上述代码复制,粘贴到YOLOv8项目的 ultralytics-main/ultralytics/nn/modules/conv.py 中;
然后将下面内容加到conv.py指定位置中,
'ChannelAttention', 'SpatialAttention', 'CBAM',
然后在下图中指定位置,加入‘CBAM’的引用。
然后,在task.py中指定位置import导入CBAM模块;
在task.py中,加入以下代码:
elif m is CBAM:
c1,c2 = ch[f],args[0]
if c2 != nc:
c2 = make_divisible(min(c2, max_channels) * width, 8)
args = [c1, *args[1:]]
至此,CBAM就成功添加到了YOLOv8中。如果想要使用,就可以在yolov8.yaml中添加使用。