空间注意力模块与通道注意力模块实现

注意力机制（CBAM）论文：1807.06521.pdf (arxiv.org)

参考博文：CBAM实现(pytorch)_cabm pytroch-CSDN博客

1.空间注意力机制（SAB）

强调感兴趣的区域同时抑制不相关的背景区域。感兴趣的区域生成一个较大的权重，不感兴趣的区域生成一个较小的权重，最后将输入特征图F乘以权重。

代码实现：

class SpatialAttention(nn.Module):
    # 空间注意力模块实现
    def __init__(self, kernel_size=7):
        super(SpatialAttention, self).__init__()

        self.conv1 = nn.Conv2d(2, 1, kernel_size, padding=kernel_size // 2, bias=False)
        self.sigmoid = nn.Sigmoid()

    def forward(self, x):
        avg_out = torch.mean(x, dim=1, keepdim=True)  # 用于计算在第一个维度上的均值，即对每一行进行求平均，keepdim=True表示保持结果的维度与输入的维度一致（即保持行数不变）
        max_out, _ = torch.max(x, dim=1, keepdim=True)
        out = torch.cat([avg_out, max_out], dim=1)  # 将每一行的均值和最大值连接在一起，形成一个新的张量
        out = self.conv1(out)
        return self.sigmoid(out) * x

2.通道注意力（CAB）

显示不同通道（一个通道是由一个卷积核对输入图像做卷积操作生成的特征图像）之间的相关性和特征图的重要程度。对不同的通道生成一个不一样的权重，权重大说明此通道比较重要。

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.fc = nn.Sequential(nn.Conv2d(in_planes, in_planes // ratio, 1, bias=False),
                                nn.ReLU(),
                                nn.Conv2d(in_planes // ratio, in_planes, 1, bias=False))  # MLP模块由两个卷积乘后再全连接
        self.sigmoid = nn.Sigmoid()

    def forward(self, x):
        avg_out = self.fc(self.avg_pool(x))
        max_out = self.fc(self.max_pool(x))
        out = avg_out + max_out  # 通道合并？
        return self.sigmoid(out) * x  # 通道注意力模块得到的权重与特征图相乘

3.CBAM注意力机制

将CAB和SAB结合起来。

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
        print('outchannels:{}'.format(out.shape))
        out = self.spatial_attention(out) * out
        return out

SE注意力机制

 

 代码实现：

class SELayer(nn.Module):
    # SE注意力机制
    def __init__(self, channel, reduction=16):
        super(SELayer, self).__init__()
        self.avg_pool = nn.AdaptiveAvgPool2d(1)  # 平局池化操作
        self.fc = nn.Sequential(
            nn.Linear(channel, channel // reduction, bias=False),  # 全连接
            nn.ReLU(inplace=True),
            nn.Linear(channel // reduction, channel, bias=False),  # 全连接
            nn.Sigmoid()
        )

    def forward(self, x):
        b, c, _, _ = x.size()  # 获取输入特征图 x 的大小，b 表示批次的大小，c 表示输入特征图的通道数
        y = self.avg_pool(x).view(b, c)
        # 输入特征图 x 经过平均池化层 self.avg_pool 处理，然后使用 view 方法将其形状变为 (b, c)。这一步是为了将特征图转换为向量的形式
        y = self.fc(y).view(b, c, 1, 1)
        # 特征向量 y 经过全连接层 self.fc 处理，然后使用 view 方法将其形状变为 (b, c, 1, 1)。这一步是为了将特征向量转换为与输入特征图相同的形状
        return x * y.expand_as(x)
        # 将输入特征图 x 与调整形状后的 y 进行逐元素相乘，得到最终的输出特征图。y.expand_as(x) 将 y 扩展为与 x 相同的形状，以便进行逐元素相乘操作