Part 2 常见的注意力机制

约 1254 字大约 4 分钟

1 SE（Squeeze and Excitation）

SE 的原始论文把研究重点放在通道维度上，探索特征通道之间的关系是否能提升网络性能。

SE 包含 Squeeze 和 Excitation 两部分。

首先是通过 Squeeze 操作 $F_{sq}$ 将输入的 $W \times H \times C$ 特征图进行全局平均池化，得到 $1 \times 1 \times C$ 的特征向量。

然后由 Excitation 操作 $F_{ex}$ 得到各通道的注意力权重，这其中又分为两步：第一步经过一个全连接层，将 $C$ 个通道变为 $C/r$ 个通道，减少参数量，然后通过 ReLU 激活函数到达第二个全连接层，将通道数恢复为 $C$ ，然后通过 Sigmoid 激活函数将权重归一化。

class SE(nn.Module):
    def __init__(self, c, r=16):
        super(SE, self).__init__()

        self.squeeze = nn.AdaptiveAvgPool2d((1, 1))

        self.excitation = nn.Sequential(
            nn.Linear(c, c // r, bias=False),
            nn.ReLU(),
            nn.Linear(c // r, c, bias=False),
            nn.Sigmoid(),
        )

    def forward(self, x):
        # 读取批次图片数量及通道数
        b, c, _, _ = x.size()

        # Squeeze 操作：经池化后输出b*c的矩阵
        y = self.squeeze(x).view(b, c)

        # Excitation 操作：经全连接层输出（b，c，1，1）矩阵
        y = self.excitation(y).view(b, c, 1, 1)

        # Scale 操作：将得到的权重乘以原来的特征图
        return x * y.expand_as(x)

2 ECA（Efficient Channel Attention）

ECA 的原始论文认为，虽然 SE 的降维可以降低模型的复杂度，但是破坏了通道与其权重之间的直接对应关系，这样权重和通道的对应关系是间接的。ECA 提出了一维卷积的方法，避免了降维对数据的影响。

ECA 首先将原始 $W \times H \times C$ 特征图进行全局平均池化得到 $1 \times 1 \times C$ 的特征向量，这一步与 SE 一致。

然后通过一个自适应卷积核 $k$ 计算各通道权重，并通过 Sigmoid 激活函数将权重归一化。

最后利用张量的广播机制，将权重加权到每个通道上。

其中自适应卷积核的大小由下式决定：

k = \left| \frac{\log_2 C}{\gamma} + \frac{b}{\gamma} \right|_{odd}

其中 $\left| \cdot \right|_{odd}$ 表示取最近的整数， $C$ 为通道数， $b=1$ ， $\gamma = 2$ 。

class ECA(nn.Module):
    def __init__(self, c, b=1, gamma=2):
        super(ECA, self).__init__()

        # 计算自适应卷积核
        kernel_size = int(abs((math.log(c, 2) + b) / gamma))
        kernel_size = kernel_size if kernel_size % 2 else kernel_size + 1

        self.gap = nn.AdaptiveAvgPool2d(1)

        self.conv = nn.Conv1d(
            1, 1, kernel_size=kernel_size, padding=(kernel_size - 1) // 2, bias=False
        )

        self.sigmoid = nn.Sigmoid()

    def forward(self, x):
        # 全局平均池化
        y = self.gap(x)

        # 1D 卷积
        y = self.conv(y.squeeze(-1).transpose(-1, -2)).transpose(-1, -2).unsqueeze(-1)

        # Sigmoid 激活
        y = self.sigmoid(y)

        return x * y.expand_as(x)

3 CBAM（Convolutional Block Attention Module）

CBAM 是一种结合了通道和空间的注意力模块，相较于 SENet 和 ECA 等只关注通道的注意力效果更佳。

首先关注通道注意力部分。该部分会对输入形状为 $H \times W \times C$ 的特征图 $F$ 分别做全局最大池化和全局平均池化，得到两个 $1 \times 1 \times C$ 的特征向量，然后经过一层 MLP 处理，得到两个新的 $1 \times 1 \times C$ 的特征向量，然后把这两个特征向量相加并使用 Sigmoid 激活函数归一化，得到通道注意力权重 $M_c$ ，其形状为 $1 \times 1 \times C$ 。

全连接层 MLP 的操作与 SE 中 Excitation 相似，都是先将 $C$ 个通道变为 $C/r$ 个通道，然后经过 ReLU 函数激活后再将通道数恢复为 $C$ 。

将通道注意力权重 $M_c$ 应用到输入特征图上 $F$ 得到特征图 $F'$ ，然后在通道维度上分别做一次全局最大池化和全局平均池化并进行拼接，得到尺寸为 $H \times W \times 2$ 的特征图。然后再做一次卷积，得到 $H \times W \times 1$ 的特征图，最后再使用 Sigmoid 激活函数归一化，得到空间注意力权重 $M_s$ 。

最后将 $M_s$ 应用在 $F'$ 上加权。

class ChannelAttention(nn.Module):
    def __init__(self, c, r=16):
        super().__init__()

        self.avg_pool = nn.AdaptiveAvgPool2d(1)
        self.max_pool = nn.AdaptiveMaxPool2d(1)

        self.mlp = nn.Sequential(
            nn.Conv2d(c, c // r, kernel_size=1, bias=False),
            nn.ReLU(inplace=True),
            nn.Conv2d(c // r, c, kernel_size=1, bias=False),
        )

        self.sigmoid = nn.Sigmoid()

    def forward(self, x):
        # 全局平均池化 + MLP
        avg_out = self.mlp(self.avg_pool(x))

        # 全局最大池化 + MLP
        max_out = self.mlp(self.max_pool(x))

        # Sigmoid 激活
        return self.sigmoid(avg_out + max_out)


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

        assert kernel_size in (3, 7), "kernel_size 必须是 3 或 7（论文默认 7）"
        padding = (kernel_size - 1) // 2
        self.conv = nn.Conv2d(2, 1, kernel_size, 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_cat = torch.cat([avg_out, max_out], dim=1)

        # 7x7 卷积
        out = self.conv(x_cat)

        # Sigmoid 激活
        return self.sigmoid(out)


class CBAM(nn.Module):
    def __init__(self, c, r=16, k=7):
        super().__init__()
        self.ca = ChannelAttention(c, r)
        self.sa = SpatialAttention(k)

    def forward(self, x):
        # 通道注意力
        ca_weight = self.ca(x)
        x = x * ca_weight
        
        # 空间注意力
        sa_weight = self.sa(x)
        x = x * sa_weight
        
        return x