# Swin Transformer 各层结构及其作用详解 Swin Transformer 是一种革命性的视觉 Transformer 架构，它通过分层设计和滑动窗口机制解决了传统 Vision Transformer 在处理高分辨率图像时计算复杂度过高的问题。下面将详细解析其各层结构及功能。 ## 整体架构概览 Swin Transformer 采用分层金字塔结构，类似于传统 CNN 的架构设计，能够生成多尺度的特征图，非常适合用于目标检测、语义分割等需要多尺度特征的计算机视觉任务 [ref_3]。 ```python # Swin Transformer 整体架构示例 import torch import torch.nn as nn class SwinTransformer(nn.Module): def __init__(self): super().__init__() # 输入处理层 self.patch_partition = PatchPartition() self.linear_embedding = LinearEmbedding() # 四个阶段（Stage）的 Transformer 块 self.stage1 = SwinTransformerBlocks(dim=96, depth=2) self.stage2 = SwinTransformerBlocks(dim=192, depth=2) self.stage3 = SwinTransformerBlocks(dim=384, depth=6) self.stage4 = SwinTransformerBlocks(dim=768, depth=2) # 下采样层 self.patch_merging_layers = nn.ModuleList([ PatchMerging(96, 192), PatchMerging(192, 384), PatchMerging(384, 768) ]) ``` ## 各层详细解析 ### 1. Patch Partition 层 **作用**：将输入图像分割成不重叠的图像块（patches），这是将图像数据转换为序列化表示的第一步 [ref_2]。 **技术细节**： - 输入：原始 RGB 图像（H × W × 3） - 输出：图像块序列（H/4 × W/4 × 48） - 通常使用 4×4 的卷积核，步长为 4 进行实现 ```python class PatchPartition(nn.Module): def __init__(self, patch_size=4): super().__init__() self.patch_size = patch_size self.conv = nn.Conv2d(3, 48, kernel_size=patch_size, stride=patch_size) def forward(self, x): # x: [B, 3, H, W] patches = self.conv(x) # [B, 48, H/4, W/4] return patches ``` ### 2. Linear Embedding 层 **作用**：将每个图像块投影到指定的特征维度，为后续的 Transformer 处理做准备 [ref_5]。 **技术细节**： - 将每个 4×4×3=48 维的图像块投影到设定的嵌入维度（如 96 维） - 相当于为每个图像块学习一个特征表示 ```python class LinearEmbedding(nn.Module): def __init__(self, in_dim=48, out_dim=96): super().__init__() self.linear = nn.Linear(in_dim, out_dim) self.layer_norm = nn.LayerNorm(out_dim) def forward(self, x): # x: [B, 48, H/4, W/4] B, C, H, W = x.shape x = x.permute(0, 2, 3, 1).reshape(B, H*W, C) # [B, N, 48] x = self.linear(x) # [B, N, 96] x = self.layer_norm(x) return x ``` ### 3. Swin Transformer Block 层 这是 Swin Transformer 的核心组件，每个阶段包含多个 Swin Transformer Block [ref_1]。 #### 3.1 W-MSA（Window Multi-head Self-Attention） **作用**：在局部窗口内计算自注意力，显著降低计算复杂度 [ref_6]。 **特点**： - 将特征图划分为不重叠的 M×M 窗口 - 只在每个窗口内部计算自注意力 - 计算复杂度从 O(n²) 降低到 O(M²×n) #### 3.2 SW-MSA（Shifted Window Multi-head Self-Attention） **作用**：通过滑动窗口机制实现跨窗口的信息交互，解决 W-MSA 的局限性 [ref_5]。 **特点**： - 在连续的块中交替使用常规窗口和滑动窗口 - 滑动窗口将特征图平移 (M/2, M/2) 个像素后重新划分窗口 - 实现不同窗口间的信息交流 ```python class SwinTransformerBlock(nn.Module): def __init__(self, dim, num_heads, window_size=7, shift_size=0): super().__init__() self.window_size = window_size self.shift_size = shift_size # 注意力机制 self.attn = WindowAttention(dim, num_heads, window_size) # MLP 层 self.mlp = nn.Sequential( nn.Linear(dim, dim * 4), nn.GELU(), nn.Linear(dim * 4, dim) ) # 归一化层 self.norm1 = nn.LayerNorm(dim) self.norm2 = nn.LayerNorm(dim) def forward(self, x): # 残差连接和层归一化 x = x + self.attn(self.norm1(x)) x = x + self.mlp(self.norm2(x)) return x ``` ### 4. Patch Merging 层 **作用**：实现特征图的下采样，构建分层特征金字塔，类似于 CNN 中的池化层 [ref_2]。 **技术细节**： - 每经过一个阶段，特征图尺寸减半，通道数加倍 - 通过拼接相邻的 2×2 图像块并线性投影实现 ```python class PatchMerging(nn.Module): def __init__(self, input_dim, output_dim): super().__init__() self.reduction = nn.Linear(4 * input_dim, output_dim) self.norm = nn.LayerNorm(output_dim) def forward(self, x): """ x: [B, H*W, C] 输出: [B, H/2*W/2, 2C] """ B, L, C = x.shape H = W = int(L ** 0.5) # 重塑为 2D 特征图 x = x.view(B, H, W, C) # 提取 2×2 邻域并拼接 x0 = x[:, 0::2, 0::2, :] # [B, H/2, W/2, C] x1 = x[:, 1::2, 0::2, :] # [B, H/2, W/2, C] x2 = x[:, 0::2, 1::2, :] # [B, H/2, W/2, C] x3 = x[:, 1::2, 1::2, :] # [B, H/2, W/2, C] x = torch.cat([x0, x1, x2, x3], -1) # [B, H/2, W/2, 4C] x = x.view(B, -1, 4 * C) # [B, H/2*W/2, 4C] x = self.reduction(x) # [B, H/2*W/2, 2C] x = self.norm(x) return x ``` ### 5. 相对位置偏置（Relative Position Bias） **作用**：在自注意力计算中引入位置信息，提升模型对空间关系的建模能力 [ref_6]。 **实现方式**： - 为每个相对位置对学习一个偏置项 - 偏置项在注意力得分计算时加到 QK^T 上 - 相对于绝对位置编码，相对位置偏置在窗口滑动时具有平移不变性 ```python class RelativePositionBias(nn.Module): def __init__(self, window_size, num_heads): super().__init__() self.window_size = window_size self.num_heads = num_heads # 相对位置偏置表 self.relative_position_bias_table = nn.Parameter( torch.zeros((2 * window_size - 1) * (2 * window_size - 1), num_heads) ) # 生成相对位置索引 coords = torch.arange(window_size) coords = torch.stack(torch.meshgrid(coords, coords)) coords_flatten = torch.flatten(coords, 1) relative_coords = coords_flatten[:, :, None] - coords_flatten[:, None, :] relative_coords = relative_coords.permute(1, 2, 0).contiguous() relative_coords[:, :, 0] += window_size - 1 relative_coords[:, :, 1] += window_size - 1 relative_coords[:, :, 0] *= 2 * window_size - 1 self.relative_position_index = relative_coords.sum(-1) def forward(self): bias = self.relative_position_bias_table[self.relative_position_index.view(-1)] bias = bias.view(self.window_size * self.window_size, self.window_size * self.window_size, -1) bias = bias.permute(2, 0, 1).contiguous() return bias ``` ## 各阶段配置与特征变化 下表展示了 Swin Transformer 四个阶段的具体配置和特征图变化： | 阶段 | Transformer 块数量 | 窗口大小 | 特征图尺寸 | 通道数 | 主要作用 | |------|-------------------|----------|------------|--------|----------| | Stage 1 | 2 | 7×7 | H/4 × W/4 | 96 | 提取局部特征，建立基础表示 | | Stage 2 | 2 | 7×7 | H/8 × W/8 | 192 | 扩大感受野，融合相邻特征 | | Stage 3 | 6 | 7×7 | H/16 × W/16 | 384 | 建立语义级特征表示 | | Stage 4 | 2 | 7×7 | H/32 × W/32 | 768 | 生成高级语义特征，用于下游任务 | ## 实际应用示例 Swin Transformer 作为骨干网络在各种视觉任务中表现出色。例如在目标检测任务中，可以将 Swin Transformer 的特征输出接入检测头 [ref_4]： ```python # YOLOX 中使用 Swin Transformer 作为骨干网络 class YOLOXWithSwin(nn.Module): def __init__(self): super().__init__() self.backbone = SwinTransformer() # 从不同阶段提取特征用于检测 self.neck = YOLOPAFPN() # 特征金字塔网络 self.head = YOLOXHead() # 检测头 def forward(self, x): # 提取多尺度特征 features = self.backbone(x) # 特征金字塔融合 neck_features = self.neck(features) # 检测预测 predictions = self.head(neck_features) return predictions ``` ## 技术优势总结 Swin Transformer 通过其分层结构设计，成功解决了传统 Vision Transformer 的多个痛点 [ref_3]： 1. **计算效率**：窗口注意力机制将计算复杂度从二次降为线性 2. **多尺度特征**：分层设计自然生成多尺度特征图 3. **平移不变性**：相对位置偏置增强了模型的空间建模能力 4. **灵活性**：适用于各种分辨率的输入图像 这种架构设计使得 Swin Transformer 在图像分类、目标检测、语义分割等多个计算机视觉任务中都取得了 state-of-the-art 的性能表现，成为当前视觉 Transformer 领域的重要里程碑。