swin transformer添加注意力机制

### 如何在Swin Transformer模型中实现或添加注意力机制 Swin Transformer 是一种基于窗口划分的分层架构，它通过局部-全局注意力机制实现了高效的视觉建模[^1]。为了增强 Swin Transformer 的表现能力，可以引入额外的注意力机制来捕捉更丰富的特征表示。 #### 1. **标准窗口注意力** Swin Transformer 已经内置了一种特殊的多头自注意力机制——窗口化多头自注意力（Window-based Multi-head Self-Attention）。这种机制将输入图像划分为固定大小的非重叠窗口，并在每个窗口内部执行自注意力操作，从而显著降低了计算复杂度[^4]。 具体而言，窗口化的多头自注意力可以通过以下方式定义： ```python class WindowAttention(nn.Module): def __init__(self, dim, window_size, num_heads, qkv_bias=True, attn_drop=0., proj_drop=0.): super().__init__() self.dim = dim self.window_size = window_size self.num_heads = num_heads # 初始化查询、键和值矩阵 self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias) # 注意力权重缩放因子 self.scale = (dim // num_heads) ** -0.5 # Dropout 层 self.attn_drop = nn.Dropout(attn_drop) self.proj = nn.Linear(dim, dim) self.proj_drop = nn.Dropout(proj_drop) def forward(self, x): B_, N, C = x.shape qkv = self.qkv(x).reshape(B_, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4) q, k, v = qkv.unbind(0) # 计算注意力分数并应用 softmax 函数 attn = (q @ k.transpose(-2, -1)) * self.scale attn = attn.softmax(dim=-1) attn = self.attn_drop(attn) # 加权求和得到最终输出 x = (attn @ v).transpose(1, 2).reshape(B_, N, C) x = self.proj(x) x = self.proj_drop(x) return x ``` --- #### 2. **扩展注意力机制** 除了默认的窗口注意力外，还可以结合其他类型的注意力机制来提升 Swin Transformer 的性能： ##### a. **SE 注意力机制** SENet（Squeeze-and-Excitation Network）是一种通道注意力机制，能够动态调整不同通道的重要性。可以在 Swin Transformer 的残差块之后加入 SE 模块，以强化重要特征的学习效果[^5]。 ```python import torch.nn as nn class SELayer(nn.Module): def __init__(self, channel, reduction=16): super(SELayer, self).__init__() self.avg_pool = nn.AdaptiveAvgPool2d(1) self.fc = nn.Sequential( nn.Linear(channel, channel // reduction, bias=False), nn.ReLU(inplace=True), nn.Linear(channel // reduction, channel, bias=False), nn.Sigmoid() ) def forward(self, x): b, c, _, _ = x.size() y = self.avg_pool(x).view(b, c) y = self.fc(y).view(b, c, 1, 1) return x * y.expand_as(x) ``` ##### b. **CBAM 注意力机制** CBAM（Convolutional Block Attention Module）综合了通道注意力和空间注意力两种形式。它可以作为附加模块嵌入到 Swin Transformer 中，进一步提高特征表达能力[^2]。 ```python 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.conv1 = 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) out = torch.cat([avg_out, max_out], dim=1) out = self.conv1(out) return self.sigmoid(out) ``` ##### c. **ECA 注意力机制** ECA（Efficient Channel Attention）是对 SE 注意力的一种优化版本，减少了参数量的同时保持了良好的性能。适合用于轻量化场景下的 Swin Transformer 改进[^3]。 ```python class ECALayer(nn.Module): """Constructs a ECA module. Args: channels: Number of input feature map's channels. gamma: Used to calculate the number of reduced dimensions `k`. Default value is set according to original paper. b: Used to calculate the number of reduced dimensions `k`. Default value is set according to original paper. """ def __init__(self, channels, gamma=2, b=1): super(ECALayer, self).__init__() t = int(abs((math.log(channels, 2) + b) / gamma)) k = t if t % 2 else t + 1 self.avg_pool = nn.AdaptiveAvgPool2d(1) self.conv = nn.Conv1d(1, 1, kernel_size=k, padding=(k - 1) // 2, bias=False) self.sigmoid = nn.Sigmoid() def forward(self, x): y = self.avg_pool(x) y = self.conv(y.squeeze(-1).transpose(-1, -2)).transpose(-1, -2).unsqueeze(-1) y = self.sigmoid(y) return x * y.expand_as(x) ``` --- #### 3. **注意事项** 当向 Swin Transformer 添加新的注意力机制时，需特别关注以下几个方面： - **维度一致性**：确保新引入的注意力模块与原始模型各层的输入/输出尺寸相匹配。 - **计算开销控制**：虽然某些注意力机制能带来性能增益，但也可能增加训练时间或内存消耗。因此应合理评估其实际效益[^4]。 - **代码结构调整**：如果是在现有框架基础上进行改造，则需要同步更新配置文件以及初始化逻辑。 ---