CBAM神经网络：从注意力机制到实战应用解析

1次阅读

没有评论

共计 3182 个字符，预计需要花费 8 分钟才能阅读完成。

在计算机视觉领域，卷积神经网络 (CNN) 长期以来一直是主流架构。然而，传统 CNN 在处理复杂场景时存在明显的局限性——它们对所有区域和通道都给予同等重视，导致模型无法专注于最重要的特征。

传统 CNN 在处理以下场景时表现欠佳：

遮挡严重的物体识别
小目标检测
复杂背景下的语义分割

这些问题的根源在于缺乏有效的特征选择机制，而注意力机制为解决这一问题提供了新思路。

CBAM(Convolutional Block Attention Module)由两个顺序子模块组成：

通道注意力模块(Channel Attention Module)
空间注意力模块(Spatial Attention Module)

这种双重注意力机制能够分别在通道和空间维度上自适应地调整特征重要性。

通道注意力的计算过程可以表示为：

$$M_c(F) = \sigma(MLP(AvgPool(F)) + MLP(MaxPool(F)))$$

其中：

$F$ 是输入特征图，维度为 $C×H×W$
$\sigma$ 是 sigmoid 激活函数
MLP 由两个全连接层组成，形成瓶颈结构

具体实现时，我们首先对特征图同时进行平均池化和最大池化，得到两个 $C×1×1$ 的向量，然后通过共享权重的 MLP 进行处理，最后将结果相加并通过 sigmoid 激活。

空间注意力的计算过程为：

$$M_s(F) = \sigma(f^{7×7}([AvgPool(F);MaxPool(F)]))$$

其中：

$[·;·]$ 表示通道拼接
$f^{7×7}$ 是 7×7 卷积核
输入是 $2×H×W$ 的特征图(平均和最大池化结果的拼接)

指标	SE-Net	CBAM
参数量	2C²/r	C²/r + 49
计算量(FLOPs)	2HWC²/r	HWC²/r + 49HW
ImageNet Top-1 Acc	+1.5%	+2.3%

注：r 是 MLP 的压缩比，通常设为 16

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)

        # 瓶颈结构的 MLP
        self.fc1 = nn.Conv2d(in_planes, in_planes // ratio, 1, bias=False)
        self.relu = nn.ReLU()
        self.fc2 = nn.Conv2d(in_planes // ratio, in_planes, 1, bias=False)

        self.sigmoid = nn.Sigmoid()

    def forward(self, x):
        # x 维度: [B, C, H, W]
        avg_out = self.fc2(self.relu(self.fc1(self.avg_pool(x))))  # [B, C, 1, 1]
        max_out = self.fc2(self.relu(self.fc1(self.max_pool(x))))  # [B, C, 1, 1]
        out = avg_out + max_out
        return self.sigmoid(out) * x  # 广播机制自动扩展

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.conv = nn.Conv2d(2, 1, kernel_size, padding=padding, bias=False)
        self.sigmoid = nn.Sigmoid()

    def forward(self, x):
        # x 维度: [B, C, H, W]
        avg_out = torch.mean(x, dim=1, keepdim=True)  # [B, 1, H, W]
        max_out, _ = torch.max(x, dim=1, keepdim=True)  # [B, 1, H, W]
        out = torch.cat([avg_out, max_out], dim=1)  # [B, 2, H, W]
        out = self.conv(out)  # [B, 1, H, W]
        return self.sigmoid(out) * x

def forward(self, x):
    residual = x

    # 原始 ResNet 块的前向传播
    out = self.conv1(x)
    out = self.bn1(out)
    out = self.relu(out)

    out = self.conv2(out)
    out = self.bn2(out)

    # 添加 CBAM 模块
    out = self.ca(out) * out  # 通道注意力
    out = self.sa(out) * out  # 空间注意力

    if self.downsample is not None:
        residual = self.downsample(x)

    out += residual
    out = self.relu(out)
    return out

当使用 DataParallel 或 DistributedDataParallel 时，CBAM 模块需要注意：