基于Transformer的CBAM模块优化实践：从注意力机制到模型性能提升

1次阅读

没有评论

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

在深度学习领域，注意力机制已成为提升模型性能的重要手段。然而，传统的注意力机制在应用中仍存在一些显著局限：

计算复杂度高：特别是全局注意力机制，其计算复杂度与输入尺寸的平方成正比，限制了在大规模输入上的应用。
信息整合不足：多数注意力机制仅关注单一维度（如通道或空间），缺乏对不同维度信息的协同建模能力。
梯度传播不稳定：深层网络中注意力模块的梯度可能变得不稳定，影响模型收敛。

CBAM（Convolutional Block Attention Module）和 Transformer 架构在注意力机制实现上具有天然的互补性：

维度互补 ：CBAM 同时建模通道和空间注意力，而 Transformer 的自注意力机制擅长捕获长程依赖关系。
计算效率 ：CBAM 的轻量级设计可以弥补 Transformer 计算开销大的问题。
特征融合 ：CBAM 的串行注意力结构能与 Transformer 的多头注意力形成层次化特征表示。

我们提出的混合架构包含三个核心组件：

CBAM-Transformer Block：将 CBAM 模块嵌入 Transformer 的前馈网络中。
跨尺度特征融合 ：在不同深度的 Transformer 层间引入 CBAM 引导的特征重组。
动态权重分配 ：通过可学习参数自动调节 CBAM 与自注意力机制的贡献比例。

注意力头数：建议 4 - 8 个头平衡计算效率和表现力
CBAM 压缩比：通道注意力中使用 16-32 的降维比例
位置编码：采用可学习的相对位置编码
归一化：LayerNorm 与 BatchNorm 的组合使用

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

class CBAM(nn.Module):
    """Convolutional Block Attention Module"""
    def __init__(self, channels, reduction=16):
        super().__init__()
        # 通道注意力
        self.avg_pool = nn.AdaptiveAvgPool2d(1)
        self.max_pool = nn.AdaptiveMaxPool2d(1)
        self.mlp = nn.Sequential(nn.Linear(channels, channels // reduction),
            nn.ReLU(),
            nn.Linear(channels // reduction, channels)
        )
        # 空间注意力
        self.conv = nn.Conv2d(2, 1, kernel_size=7, padding=3)

    def forward(self, x):
        # 通道注意力
        avg_out = self.mlp(self.avg_pool(x).squeeze(-1).squeeze(-1))
        max_out = self.mlp(self.max_pool(x).squeeze(-1).squeeze(-1))
        channel_att = torch.sigmoid(avg_out + max_out).unsqueeze(-1).unsqueeze(-1)
        x = x * channel_att

        # 空间注意力
        avg_out = torch.mean(x, dim=1, keepdim=True)
        max_out = torch.max(x, dim=1, keepdim=True)[0]
        spatial_att = torch.cat([avg_out, max_out], dim=1)
        spatial_att = torch.sigmoid(self.conv(spatial_att))
        return x * spatial_att

class TransformerCBAMBlock(nn.Module):
    """集成 CBAM 的 Transformer Block"""
    def __init__(self, d_model, nhead, dim_feedforward=2048, dropout=0.1):
        super().__init__()
        self.self_attn = nn.MultiheadAttention(d_model, nhead, dropout=dropout)
        self.cbam = CBAM(d_model)
        self.linear1 = nn.Linear(d_model, dim_feedforward)
        self.linear2 = nn.Linear(dim_feedforward, d_model)
        self.norm1 = nn.LayerNorm(d_model)
        self.norm2 = nn.LayerNorm(d_model)
        self.dropout = nn.Dropout(dropout)

    def forward(self, src):
        # 自注意力
        src2 = self.self_attn(src, src, src)[0]
        src = src + self.dropout(src2)
        src = self.norm1(src)

        # CBAM 处理（需调整维度）b, l, c = src.shape
        src_cbam = src.transpose(1,2).view(b, c, int(l**0.5), int(l**0.5))
        src_cbam = self.cbam(src_cbam)
        src_cbam = src_cbam.view(b, c, l).transpose(1,2)

        # 前馈网络
        src2 = self.linear2(F.relu(self.linear1(src_cbam)))
        src = src + self.dropout(src2)
        src = self.norm2(src)
        return src

在 ImageNet-1K 数据集上的实验结果：