CBAM(Convolutional Block Attention Module)是一种用于计算机视觉任务的注意力机制,旨在增强深度神经网络在空间和通道维度上的特征表示能力。
CBAM模块包含两个关键组件:通道注意力模块(Channel Attention Module)和空间注意力模块(Spatial Attention Module)。它们分别在通道维度和空间维度上学习特征的重要性,并将两者结合以生成最终的注意力权重。CBAM模块可以被集成到卷积神经网络的不同层级,以帮助网络更好地捕获重要特征。
以下是CBAM模块的主要步骤:
1、通道注意力模块(Channel Attention Module):通道注意力模块主要关注特征图中的通道信息,以学习每个通道的重要性权重。通道注意力模块包含以下几个步骤:
a. 全局平均池化(Global Average Pooling):将特征图的每个通道内的元素取平均,得到每个通道的平均值。这样得到的向量可以反映每个通道的全局特征。
b. 全连接层(Fully Connected Layer):使用一个全连接层将全局平均池化得到的向量映射到一个中间维度。在这个全连接层中,通常使用ReLU激活函数来引入非线性。
c. 全连接层(Fully Connected Layer):再使用一个全连接层将中间维度的向量映射回原始通道维度,生成注意力权重。在这个全连接层中,通常使用Sigmoid激活函数,将注意力权重限制在0到1之间。
# 通道注意力模块
class ChannelAttentionModule(nn.Module):
def __init__(self, channel, reduction=16):
super(ChannelAttentionModule, self).__init__()
mid_channel = channel // reduction
# 使用自适应池化缩减map的大小,保持通道不变
self.avg_pool = nn.AdaptiveAvgPool2d(1)
self.max_pool = nn.AdaptiveMaxPool2d(1)
self.shared_MLP = nn.Sequential(
nn.Linear(in_features=channel, out_features=mid_channel),
nn.ReLU(),
nn.Linear(in_features=mid_channel, out_features=channel)
)
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 ChannelAttentionModule(nn.Module):
- 这是一个Python类定义,它继承自
nn.Module
,这是PyTorch中用于构建神经网络的基类。
def __init__(self, channel, reduction=16):
- 这是类的构造函数(
__init__
方法),用于初始化类的实例。它接受两个参数:channel
和reduction
。 channel
: 表示输入张量的通道数,也可以理解为特征图的深度。reduction
: 是一个可选参数,默认值为16,用于控制通道的缩减比例。
super(ChannelAttentionModule, self).__init__()
- 这一行调用父类
nn.Module
的构造函数,以确保正确地初始化继承自nn.Module
的功能。
mid_channel = channel // reduction
- 这一行计算
mid_channel
,它表示在通道缩减后的中间通道数。这里使用了整数除法运算符//
,保证mid_channel
是一个整数。
self.avg_pool = nn.AdaptiveAvgPool2d(1)
- 这一行创建一个自适应平均池化层,用于对输入特征图进行平均池化操作。
nn.AdaptiveAvgPool2d
会将输入的任意大小的特征图转换为大小为1x1的特征图,同时保留通道数不变。
self.max_pool = nn.AdaptiveAvgPool2d(1)
- 这一行创建一个自适应最大池化层,与平均池化层类似,但它会对输入特征图进行最大池化操作。
self.shared_MLP = nn.Sequential(...)
- 这一行定义了一个包含两个线性层和一个激活函数ReLU的多层感知机(MLP)模型。这个MLP是通道注意力模块的核心部分。
- 第一个线性层将输入的通道数
channel
缩减到mid_channel
,然后通过ReLU激活函数。 - 第二个线性层将
mid_channel
的特征再扩展回原始的通道数channel
,用于恢复通道维度的信息。
self.sigmoid = nn.Sigmoid()
- 这一行创建了一个Sigmoid激活函数,用于将输出值映射到0到1之间的概率范围。
def forward(self, x):
- 这是类的前向传播函数(
forward
方法),用于定义数据在模型中的传递方式。
avgout = self.shared_MLP(self.avg_pool(x).view(x.size(0), 1)).unsqueeze(2).unsqueeze(3)
- 这一行计算了平均池化分支的输出。首先,输入特征图
x
经过平均池化层得到大小为1x1的特征图。接着,利用view
方法将通道维度的特征转换为一维,以适应MLP的输入。然后,通过self.shared_MLP
进行处理,并使用unsqueeze
方法扩展维度,以便后续的操作。
maxout = self.shared_MLP(self.avg_pool(x).view(x.size(0), 1)).unsqueeze(2).unsqueeze(3)
- 这一行计算了最大池化分支的输出。与前面相同,首先,输入特征图
x
经过最大池化层得到大小为1x1的特征图。然后,通过self.shared_MLP
进行处理,并使用unsqueeze
方法扩展维度。
return self.sigmoid(avgout + maxout)
- 最后,将平均池化分支和最大池化分支的结果相加,并经过Sigmoid激活函数,得到通道注意力模块的输出。这将会产生一个大小为1x1x1x1的特征图,其中的值表示每个通道的重要性权重,值越接近1表示该通道对应的特征越重要,越接近0表示该通道对应的特征不重要。
2、空间注意力模块(Spatial Attention Module):空间注意力模块主要关注特征图中的空间信息,以学习每个空间位置的重要性权重。空间注意力模块包含以下几个步骤:
a. 最大池化和平均池化(Max Pooling and Average Pooling):分别对特征图进行最大池化和平均池化操作,得到两个不同的空间信息向量。
b. 堆叠(Concatenation):将最大池化和平均池化得到的两个向量进行堆叠,形成一个新的向量。这个向量蕴含了每个空间位置的不同池化信息。
c. 卷积层(Convolutional Layer):使用一个卷积层对堆叠后的向量进行特征融合和映射。
d. 全连接层(Fully Connected Layer):使用一个全连接层将卷积层的输出映射到一个中间维度。在这个全连接层中,通常使用ReLU激活函数。
e. 全连接层(Fully Connected Layer):再使用一个全连接层将中间维度的向量映射回原始的空间维度,生成注意力权重。在这个全连接层中,通常使用Sigmoid激活函数,将注意力权重限制在0到1之间。
# 空间注意力模块
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):
# map尺寸不变,缩减通道
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
self.conv2d = nn.Conv2d(in_channels=2, out_channels=1, kernel_size=7, stride=1, padding=3)
-
定义一个 2D 卷积层
conv2d
,用于学习空间注意力权重。这个卷积层的输入通道数为 2,输出通道数为 1,卷积核大小为 7x7,步长为 1,填充为 3。
self.sigmoid = nn.Sigmoid()
-
创建一个 Sigmoid 激活函数,用于将输出限制在 0 到 1 之间,以得到注意力权重。
def forward(self, x):
- 定义了空间注意力模块的前向传播过程,输入为特征图
x
avgout = torch.mean(x, dim=1, keepdim=True)
- 计算输入特征图
x
在通道维度上的平均值,即对每个空间位置上的通道值取平均。得到的avgout
形状为(batch_size, 1, height, width)
,其中height
和width
是特征图的高度和宽度。
maxout, _ = torch.max(x, dim=1, keepdim=True)
- 计算输入特征图
x
在通道维度上的最大值,即对每个空间位置上的通道值取最大。得到的maxout
形状同样为(batch_size, 1, height, width)。
out = torch.cat([avgout, maxout], dim=1)
- 将平均值
avgout
和最大值maxout
在通道维度上进行拼接,得到一个新的特征图out
,其中通道数为 2。
out = self.sigmoid(self.conv2d(out))
- 通过 2D 卷积层和 Sigmoid 激活函数计算空间注意力权重。卷积操作将学习每个空间位置的权重,Sigmoid 激活函数将输出限制在 0 到 1 之间,表示每个空间位置的重要性。
return out
- 最终输出经过空间注意力模块处理后的特征图
out
,其中每个空间位置都有一个对应的注意力权重。
3、综合注意力权重:将通道注意力模块和空间注意力模块得到的注意力权重相乘,得到最终的综合注意力权重。
4、特征加权:使用综合的注意力权重对输入特征图进行加权求和,得到加权后的特征表示。
# CBAM模块
class CBAM(nn.Module):
def __init__(self, channel):
super(CBAM, self).__init__()
self.channel_attention = ChannelAttentionModule(channel)
self.spatial_attention = SpatialAttentionModule()
def forward(self, x):
out = self.channel_attention(x) * x
out = self.spatial_attention(out) * out
return out
def __init__(self, channel):
- 这是 CBAM 类的构造函数,
channel
表示输入特征图的通道数。在初始化过程中,创建了通道注意力模块和空间注意力模块。
self.channel_attention = ChannelAttentionModule(channel)
- 在构造函数中,创建了通道注意力模块
channel_attention
,并传入输入特征图的通道数channel
self.spatial_attention = SpatialAttentionModule()
- 在构造函数中,创建了空间注意力模块
spatial_attention
def forward(self, x):
- 这是 CBAM 类的前向传播函数,定义了 CBAM 模块的前向计算过程,输入为特征图
x
out = self.channel_attention(x) * x
- 通过通道注意力模块处理输入特征图
x
,得到加权后的特征图out
。通道注意力模块将学习每个通道的重要性,并将注意力权重应用到相应通道的特征图上。
out = self.spatial_attention(out) * out
- 通过空间注意力模块处理上一步得到的特征图
out
,得到最终加权后的特征图out
。空间注意力模块将学习每个空间位置的重要性,并将注意力权重应用到相应位置的特征图上。
return out
- 返回经过 CBAM 模块处理后的特征图
out
,其中每个通道和空间位置都有相应的注意力权重被应用。这样的处理有助于提升深度神经网络对重要特征的捕捉和表示能力。
通过以上步骤,CBAM模块能够同时考虑特征图中每个通道和每个空间位置的重要性,从而增强网络对重要特征的表示和选择,提高模型性能。