【论文笔记】CBAM：Convolutional Block Attention Module（卷积块注意力模型）

本文是关于论文《CBAM: Convolutional Block Attention Module》的阅读笔记。

一、概述

深度学习中的注意力机制是模仿人类的注意力产生的，当我们在看一副图像时，图像中有些地方（比如人物照中的人的面部）会引起我们的注意力，而其他地方则可能会忽视。这相当于给整幅图像加了一个可视化的权重。注意力机制可以分为通道注意力和空间注意力等。简单来说，注意力机制就是给每个通道或整个空间特征图乘以一个权重，以表示它们的重要性。

CBAM包含两部分，一是通道的注意力，二是空间的注意力。CBAM可以用以下公式表示：
$$\begin{array}{rl}{\mathbf{F}}^{\prime }& ={\mathbf{M}}_{\mathbf{c}}(\mathbf{F})\otimes \mathbf{F}\\ {\mathbf{F}}^{\prime \prime }& ={\mathbf{M}}_{\mathbf{s}}\left({\mathbf{F}}^{\prime }\right)\otimes {\mathbf{F}}^{\prime }\end{array}$$
其中，$F$是特征图，$M_c$和$M_s$分别表示基于通道的和基于空间的注意力，$\otimes$表示逐元素相乘，$F^{\prime }$和$F^{\prime \prime }$分别表示进行了通道注意力和空间注意力后的输出特征图。由于CBAM模块的输入和输出大小相同，所以它可以插入到已有模型的任意位置。

二、通道注意力模块

通道注意力模块先对输入特征图在空间维度分别进行最大池化和平均池化，然后分别经过一个共享权重的MLP（多层感知机） ，然后将两者的输出做逐元素的相加，再经过sigmoid激活函数，得到通道注意力权重，将该权重和输入特征图做逐元素的乘法，就实现了通道上的注意力机制。

在空间维度分别进行最大池化和平均池化，其实就是只保留通道的维度，其他维度为1，得到一个一维的矢量，然后再经过共享权重的MLP，MLP由两层卷积操作组成，第一个卷积后使用ReLU激活函数，第二个卷积后使用sigmoid激活函数得到输出权重，然后将每个通道乘以其对应的权重。上述过程可以表示为：
$$\begin{array}{rl}{\mathbf{M}}_{\mathbf{c}}(\mathbf{F})& =\sigma (\mathrm{MLP}(A\mathrm{vg}\mathrm{Pool}(\mathbf{F}))+MLP(\mathrm{Max}\mathrm{Pool}(\mathbf{F})))\\ & =\sigma \left({\mathbf{W}}_1\left({\mathbf{W}}_0\left({\mathbf{F}}_{\mathbf{a}\mathbf{v}\mathbf{g}}^{\mathbf{c}}\right)\right)+{\mathbf{W}}_1\left({\mathbf{W}}_0\left({\mathbf{F}}_{\mathbf{m}\mathbf{a}\mathbf{x}}^{\mathbf{c}}\right)\right)\right)\end{array}$$
其中$\sigma$表示sigmoid激活函数，$W_0$和$W_1$分别表示两个卷积操作，$F_{avg}^c$和$F_{max}^c$分别表示平均池化和最大池化。

三、空间注意力模块

先对输入特征图做基于通道的最大池化和平均池化操作，将得到的两个特征图做基于通道的拼接，再对其进行卷积操作和sigmoid激活函数，得到通道数为1，特征图大小和输入相同的注意力权重，再与输入特征图做逐元素的乘法，就完成了空间注意力机制。该过程可以表示为：
$$\begin{array}{rl}{\mathbf{M}}_{\mathbf{s}}(\mathbf{F})& =\sigma \left(f^{7\times 7}([\text{AvgPool}(\mathbf{F});\mathrm{Max}\mathrm{Pool}(\mathbf{F})])\right)\\ & =\sigma \left(f^{7\times 7}\left({\mathbf{F}}_{\mathbf{a}\mathbf{v}\mathbf{g}}^{\mathbf{s}};{\mathbf{F}}_{\mathbf{m}\mathbf{a}\mathbf{x}}^{\mathbf{s}}\right]\right))\end{array}$$
其中$f^{7\times 7}$表示卷积核大小为$7\times 7$的卷积操作。

四、代码

import torch
import torch.nn as nn
import torch.nn.functional as F


class CBAM_Module(nn.Module):
    def __init__(self, dim, in_channels, ratio, kernel_size):
        super(CBAM_Module, self).__init__()
        self.avg_pool = getattr(nn, "AdaptiveAvgPool{0}d".format(dim))(1)
        self.max_pool = getattr(nn, "AdaptiveMaxPool{0}d".format(dim))(1)
        conv_fn = getattr(nn, "Conv{0}d".format(dim))
        self.fc1 = conv_fn(in_channels, in_channels // ratio, kernel_size=1, padding=0)
        self.relu = nn.ReLU()
        self.fc2 = conv_fn(in_channels // ratio, in_channels, kernel_size=1, padding=0)
        self.sigmoid = nn.Sigmoid()
        assert kernel_size in (3, 7), 'kernel size must be 3 or 7'
        padding = 3 if kernel_size == 7 else 1
        self.conv = conv_fn(2, 1, kernel_size=kernel_size, stride=1, padding=padding)

    def forward(self, x):
        # Channel attention module:（Mc(f) = σ(MLP(AvgPool(f)) + MLP(MaxPool(f)))）
        module_input = x
        avg = self.fc2(self.relu(self.fc1(self.avg_pool(x))))
        mx = self.fc2(self.relu(self.fc1(self.max_pool(x))))
        x = self.sigmoid(avg + mx)
        x = module_input * x
        # Spatial attention module:Ms (f) = σ( f7×7( AvgPool(f) ; MaxPool(F)] )))
        module_input = x
        avg = torch.mean(x, dim=1, keepdim=True)
        mx, _ = torch.max(x, dim=1, keepdim=True)
        x = torch.cat((avg, mx), dim=1)
        x = self.sigmoid(self.conv(x))
        x = module_input * x
        return x