CBAM注意力机制结构图自绘版：从原理到实现的全方位解析

1次阅读

没有评论

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

在传统的 CNN 模型中，所有特征通道和空间位置通常被平等对待。但事实上，不同通道和区域对最终结果的贡献度差异显著。例如：

通道维度：某些特征通道可能包含更多噪声
空间维度：图像的关键信息往往集中在特定区域

传统方法通过增加网络深度或宽度来提升性能，但这种粗放的方式会导致：

计算资源浪费在无关特征上
模型难以自适应聚焦关键区域
特征利用率低下

CBAM 包含两个串行的子模块：

graph LR
    A[输入特征] --> B[通道注意力] --> C[空间注意力] --> D[输出特征]

关键计算步骤：

分别进行全局平均池化和最大池化：
$$\text{AvgPool}(F) = \frac{1}{H\times W}\sum_{i=1}^H\sum_{j=1}^W F_{i,j}$$
$$\text{MaxPool}(F) = \max_{i,j} F_{i,j}$$
共享 MLP 网络处理（含降维）：
$$W_0 \in \mathbb{R}^{C/r \times C}, W_1 \in \mathbb{R}^{C \times C/r}$$
合并两种池化结果：
$$M_c(F) = \sigma(\text{MLP}(\text{AvgPool}(F)) + \text{MLP}(\text{MaxPool}(F)))$$

实现步骤：

沿通道维进行两种池化：
$$\text{AvgPool}(F) \in \mathbb{R}^{H\times W}, \text{MaxPool}(F) \in \mathbb{R}^{H\times W}$$
拼接后卷积处理：
$$M_s(F) = \sigma(f^{7\times7}([\text{AvgPool}(F); \text{MaxPool}(F)]))$$

import torch
import torch.nn as nn

class ChannelAttention(nn.Module):
    def __init__(self, in_planes, ratio=16):
        super().__init__()
        self.avg_pool = nn.AdaptiveAvgPool2d(1)
        self.max_pool = nn.AdaptiveMaxPool2d(1)

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

class SpatialAttention(nn.Module):
    def __init__(self, kernel_size=7):
        super().__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):
        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.conv(x)
        return self.sigmoid(x)

class CBAM(nn.Module):
    def __init__(self, channels, ratio=16, kernel_size=7):
        super().__init__()
        self.ca = ChannelAttention(channels, ratio)
        self.sa = SpatialAttention(kernel_size)

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

在 CIFAR-10 上的测试结果（ResNet18 backbone）：