[论文阅读]CBAM——代码实现和讲解

news2024/11/26 8:23:13

CBAM

论文网址：CBAM
论文代码：CBAM

CBAM
本文提出了一种卷积块注意力模块（CBAM），它是卷积神经网络（CNN）的一种轻量级、高效的注意力模块。该模块沿着通道和空间两个独立维度依次推导注意力图，然后将这些注意力图与输入特征图相乘，进行自适应特征细化。CBAM 可以无缝集成到任何 CNN 架构中，开销几乎可以忽略不计，并且可以与基本 CNN 一起进行端到端训练。作者通过在 ImageNet-1K、MS COCO 检测和 VOC 2007 检测数据集上的大量实验验证了 CBAM，结果表明各种模型在分类和检测性能上都有一致的提高。论文还讨论了网络工程和注意力机制方面的相关工作，并强调了所提出的 CBAM 模块的贡献。

CBAM 模块由两个分支组成：通道注意力分支和空间注意力分支。

Channer Attention Module

通道注意力分支旨在捕捉特征图不同通道之间的相互依存关系。它使用全局平均池化操作来获取通道统计信息，然后应用两个全连接层来生成通道注意图。然后将这些注意力图与原始特征图按要素相乘，以突出重要通道。
Channer Attention Module

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 // 16, 1, bias=False),
                               nn.ReLU(),
                               nn.Conv2d(in_planes // 16, in_planes, 1, bias=False))
        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)

Spatial Attention Module

另一方面，空间注意力分支侧重于捕捉每个通道内的空间依赖性。它采用挤压-激发操作，首先使用全局平均池化技术对输入特征图进行挤压，以获得通道统计信息。然后，应用两个全连接层生成空间注意力图。这些图随后与原始特征图按元素相乘，以强调重要的空间位置。
Spatial Attention Module

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)
        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)