通道注意力和空间注意力（CBAM）-526互联

用实例说明通道注意力机制与空间注意力机制的内容。包含代码示例。

参考的博文：

pytorch中加入注意力机制（CBAM），以ResNet为例。解析到底要不要用ImageNet预训练？如何加预训练参数？

（六十一）通俗易懂理解——通道注意力机制和空间注意力机制(CBAM)

CBAM

Convolutional Block Attention Module (CBAM) 表示卷积模块的注意力机制模块，是一种结合了空间和通道的注意力机制模块^[1]。

通道注意力机制

例如对 batch x 7 x 7 x 512 的张量使用通道注意力机制，会获得 batch x 512 的张量作为这 512 个通道的权重。

即，通道注意力机制是使网络获得对通道的权重。

具体来说，若输入张量为 $x$，则运算为：

$$
\text{softmax}(\text{Linear}(\text{AvgPool}(x))+\text{Linear}(\text{MaxPool}(x)))
$$

实现代码：

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

        self.sigmoid = nn.Sigmoid()

    def forward(self, x):
        avg_out = self.fc2(self.relu1(self.fc1(self.avg_pool(x))))
        max_out = self.fc2(self.relu1(self.fc1(self.max_pool(x))))
        out = avg_out + max_out
        return self.sigmoid(out)

空间注意力机制

例如对 batch x 7 x 7 x 512 的张量使用通道注意力机制，会获得 batch x 7 x 7 的张量作为这 7 x 7 空间的权重。

实现代码：

class SpatialAttention(nn.Module):
    def __init__(self, kernel_size=7):
        super(SpatialAttention, self).__init__()

        assert kernel_size in (3, 7), 'kernel size must be 3 or 7'
        padding = 3 if kernel_size == 7 else 1

        self.conv1 = nn.Conv2d(2, 1, kernel_size, padding=padding, 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)

ResNet 添加注意力机制

以 Resnet 为例。

例如输入 Resnet 网络的图像大小为 224 x 224，网络最后会生成一个维度在 batch x 7 x 7 x 512 的张量，然后经过 nn.AdaptiveAvgPool2d((1, 1)) 与 nn.Linear 获得维度在 batch x 1000 的张量。

在送入 nn.AdaptiveAvgPool2d((1, 1)) 与 nn.Linear 层之前，进行通道注意力和空间注意力。

class ResNet(nn.Module):

    def __init__():
        ···
        self.ca = ChannelAttention(self.inplanes)
        self.sa = SpatialAttention()

    def forward(self, x):
        ···

        x = self.layer4(x)

        x = self.ca(x) * x
        x = self.sa(x) * x

        ···

        return x

Woo, Sanghyun, et al. "Cbam: Convolutional block attention module." Proceedings of the European conference on computer vision (ECCV). 2018. https://openaccess.thecvf.com/content_ECCV_2018/papers/Sanghyun_Woo_Convolutional_Block_Attention_ECCV_2018_paper.pdf ↩︎