CBAM: Convolutional Block Attention Module

 1、引言

        论文链接：CBAM: Convolutional Block Attention Module | SpringerLink

        Sanghyun Woo[1] 等将注意力过程分解为互相独立的通道注意力和空间注意力而得到CBAM(convolutional block attention module)[1]。这样不仅可以节约参数和计算量，而且使 CBAM 能无缝地集成到现有的网络架构中去，并可以和被集成网络一起端到端训练。CBAM 对各种模型的分类和检测性能都有一致的改进，即 CBAM 的适用性非常广泛，是混合注意力机制的代表模型。

2、方法

图1  CBAM

        CBAM 依次由通道注意力模块和空间注意力模块组成。如图 1 所示，对于输入特征 F，先使用通道注意力模块得到的通道权重 Mc 调整 F 的通道特征（F*Mc）得到 F'，再将 F' 输入空间注意力模块得到空间权重 Ms，使用 Ms 调整 F' 的空间特征（F‘*Ms）得到 F''，F'' 就是 CBAM 的输出。

图2  通道注意力模块

        CBAM 中的通道注意力模块可以看作是在 SE(Squeeze and Excitation) block[2] 的基础上增加了一个 AdaptiveMaxPool 分支，只是这两个分支共享一个 MLP(multiple layer perceptronz)。即对于 shape=(C, H, W) 的输入特征 F，分别使用 AdaptiveAvgPool 和 AdaptiveMaxPool 各得到一个 shape=(C, 1, 1) 的特征图；再将 AdaptiveAvgPool 和 AdaptiveMaxPool 的结果分别送入一个共享的只有一个隐藏层的 MLP 中学习，各得到 1 个 shape=(C, 1, 1) 的特征，这里使用的 MLP 和 SE block 中的完全相同；最后对这 2 个 MLP 的输出结果执行 Add 操作，然后使用 Sigmoid 激活得到 F 的通道权重 Mc。如图 2 所示。

图3  空间注意力模块

        CBAM 中的空间注意力模块如图 3 所示。对于 shape=(C, H, W) 的输入特征 F’（在 CBAM 中 F’ 是使用通道权重 Mc 调整后的特征），先分别进行通道维度的全局最大池化和全局平均池化各得到一个 shape=(1, H, W) 的特征；再将全局最大池化和全局平均池化的结果沿通道维度拼接得到一个 shape=(2, H, W) 的特征；最后经过一个 kernel_size=(7, 7) 的卷积层得到一个 shape=(1, H, W) 的特征，然后使用 Sigmoid 激活得到 F‘ 的空间权重 Ms。

图4  ResBlock+CBAM

        CBAM 和某个网络的结合方式和 SE block 完全一样，例如将 ResNet[3] 和 CBAM 结合时，CBAM 位于每个 ResBlock 的最后一个卷积层（实际上是最后一个 BatchNorm）的后面，CBAM 的输出加上恒等映射就是 ResBlock+CBAM 的输出，具体如图 4 所示。

3、总结

        [1] 提出了能提高网络表示能力的 CBAM，CBAM 依次由通道注意力模块和空间注意力模块组成，能在保持开销很小的同时实现可观的性能改进。通道注意力模块使用了最大池化特征和平均池化特征，从而能产生比 SE block 更精细的注意力，空间注意力的引入进一步提高了性能。CBAM 在 ImageNet-1K、MS COCO 和 VOC 2007 上优于所有 baseline。

        由于作者没有公开代码，故本人基于 Pytorch 实现了 CBAM 如下所示：

import torch
from torch import nn
from torch.nn import init

__all__ = ['CBAMBlock', 'ChannelAttention', 'SpatialAttention']


class ChannelAttention(nn.Module):
    def __init__(self, channel, reduction=16):
        super().__init__()
        self.maxpool = nn.AdaptiveMaxPool2d(1)
        self.avgpool = nn.AdaptiveAvgPool2d(1)
        self.se = nn.Sequential(
            nn.Conv2d(channel, channel // reduction, 1, bias=False),
            nn.ReLU(),
            nn.Conv2d(channel // reduction, channel, 1, bias=False)
        )
        self.sigmoid = nn.Sigmoid()

    def forward(self, x):
        max_result = self.maxpool(x)
        avg_result = self.avgpool(x)
        max_out = self.se(max_result)
        avg_out = self.se(avg_result)
        w = self.sigmoid(max_out + avg_out)
        return x * w


class SpatialAttention(nn.Module):
    def __init__(self, kernel_size=7):
        super().__init__()
        self.conv = nn.Conv2d(2, 1, kernel_size=kernel_size, padding=kernel_size // 2)
        self.sigmoid = nn.Sigmoid()

    def forward(self, x):
        max_result, _ = torch.max(x, dim=1, keepdim=True)
        avg_result = torch.mean(x, dim=1, keepdim=True)
        result = torch.cat([max_result, avg_result], 1)
        w = self.conv(result)
        w = self.sigmoid(w)
        return x * w


class CBAMBlock(nn.Module):

    def __init__(self, channel, reduction=16, kernel_size=7):
        super().__init__()
        self.ca = ChannelAttention(channel=channel, reduction=reduction)
        self.sa = SpatialAttention(kernel_size=kernel_size)

    def init_weights(self):
        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                init.kaiming_normal_(m.weight, mode='fan_out')
                if m.bias is not None:
                    init.constant_(m.bias, 0)

    def forward(self, x):
        out = self.ca(x)
        out = self.sa(out)
        return out

参考文献

[1] Sanghyun Woo, Jongchan Park, Joon-Young Lee, and In So Kweon. CBAM: Convolutional Block Attention Module. In ECCV, 2018.

[2] Jie Hu, Li Shen, and Gang Sun. Squeeze-and-Excitation Networks. In CVPR, 2018.

[2] Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep Residual Learning for Image Recognition. In CVPR, 2016.