别再死记硬背Swin-Transformer结构了！用PyTorch手撕W-MSA和SW-MSA，理解偏移窗口的代码实现

从零实现Swin-Transformer核心模块：W-MSA与SW-MSA的PyTorch实战解析

在计算机视觉领域，Transformer架构正在重塑传统卷积神经网络的统治地位。而Swin-Transformer作为其中的佼佼者，通过引入窗口多头自注意力（W-MSA）和移位窗口多头自注意力（SW-MSA）机制，在保持模型性能的同时大幅降低了计算复杂度。本文将带您深入这两个核心模块的实现细节，用PyTorch从零开始构建完整的功能模块。

1. 理解W-MSA的设计动机

传统视觉Transformer（ViT）在处理高分辨率图像时面临计算量平方级增长的问题。假设输入特征图尺寸为h×w，标准MSA的计算复杂度为：

Ω(MSA) = 4hwC² + 2(hw)²C

而W-MSA通过将特征图划分为不重叠的M×M窗口，仅在窗口内计算自注意力，将复杂度降低为：

Ω(W-MSA) = 4hwC² + 2M²hwC

关键优势对比：

实现窗口划分的PyTorch代码如下：

def window_partition(x, window_size): """ Args: x: (B, H, W, C) window_size (int): 窗口大小M Returns: windows: (num_windows*B, window_size, window_size, C) """ B, H, W, C = x.shape x = x.view(B, H // window_size, window_size, W // window_size, window_size, C) windows = x.permute(0, 1, 3, 2, 4, 5).contiguous().view(-1, window_size, window_size, C) return windows

2. W-MSA的完整实现

在完成窗口划分后，我们需要在每个窗口内实现标准的自注意力计算。以下是多头自注意力的关键步骤实现：

class WindowAttention(nn.Module): def __init__(self, dim, window_size, num_heads): super().__init__() self.dim = dim self.window_size = window_size self.num_heads = num_heads head_dim = dim // num_heads self.scale = head_dim ** -0.5 # 相对位置偏置表 self.relative_position_bias_table = nn.Parameter( torch.zeros((2 * window_size[0] - 1) * (2 * window_size[1] - 1), num_heads)) # 生成相对位置索引 coords_h = torch.arange(self.window_size[0]) coords_w = torch.arange(self.window_size[1]) coords = torch.stack(torch.meshgrid([coords_h, coords_w])) # 2, Wh, Ww coords_flatten = torch.flatten(coords, 1) # 2, Wh*Ww relative_coords = coords_flatten[:, :, None] - coords_flatten[:, None, :] # 2, Wh*Ww, Wh*Ww relative_coords = relative_coords.permute(1, 2, 0).contiguous() # Wh*Ww, Wh*Ww, 2 relative_coords[:, :, 0] += self.window_size[0] - 1 # 转换为非负 relative_coords[:, :, 1] += self.window_size[1] - 1 relative_coords[:, :, 0] *= 2 * self.window_size[1] - 1 relative_position_index = relative_coords.sum(-1) # Wh*Ww, Wh*Ww self.register_buffer("relative_position_index", relative_position_index) self.qkv = nn.Linear(dim, dim * 3) self.proj = nn.Linear(dim, dim) def forward(self, x, mask=None): 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[0], qkv[1], qkv[2] # 每个形状为 (B_, num_heads, N, head_dim) attn = (q @ k.transpose(-2, -1)) * self.scale # 添加相对位置偏置 relative_position_bias = self.relative_position_bias_table[ self.relative_position_index.view(-1)].view( self.window_size[0] * self.window_size[1], self.window_size[0] * self.window_size[1], -1) # Wh*Ww, Wh*Ww, nH relative_position_bias = relative_position_bias.permute(2, 0, 1).contiguous() # nH, Wh*Ww, Wh*Ww attn = attn + relative_position_bias.unsqueeze(0) if mask is not None: nW = mask.shape[0] attn = attn.view(B_ // nW, nW, self.num_heads, N, N) + mask.unsqueeze(1).unsqueeze(0) attn = attn.view(-1, self.num_heads, N, N) attn = attn.softmax(dim=-1) x = (attn @ v).transpose(1, 2).reshape(B_, N, C) x = self.proj(x) return x

关键实现细节：

1. 相对位置偏置的编码方式将二维位置关系映射到一维索引
2. 使用独立的可学习参数表relative_position_bias_table存储偏置
3. 注意力分数的计算采用缩放点积注意力机制
4. 支持可选的注意力掩码（用于SW-MSA）

3. 实现SW-MSA的窗口移位机制

SW-MSA通过引入窗口移位解决了W-MSA缺乏跨窗口信息交互的问题。其核心操作包括：

1. 对特征图进行循环移位（cyclic shift）
2. 在移位后的特征图上应用W-MSA
3. 使用注意力掩码确保只计算有效区域的自注意力
4. 反向循环移位恢复原始位置

移位操作可视化：

原始窗口划分： +---+---+---+---+ | 1 | 1 | 2 | 2 | +---+---+---+---+ | 1 | 1 | 2 | 2 | +---+---+---+---+ | 3 | 3 | 4 | 4 | +---+---+---+---+ | 3 | 3 | 4 | 4 | +---+---+---+---+ 移位后窗口划分（M=2）： +---+---+---+---+ | 4 | 1 | 1 | 2 | +---+---+---+---+ | 3 | 3 | 4 | 1 | +---+---+---+---+ | 3 | 2 | 2 | 4 | +---+---+---+---+ | 2 | 3 | 4 | 4 | +---+---+---+---+

实现移位和掩码生成的代码如下：

def create_mask(window_size, shift_size, H, W): # 计算掩码，确保只计算有效区域的自注意力 img_mask = torch.zeros((1, H, W, 1)) # 1 H W 1 h_slices = (slice(0, -window_size), slice(-window_size, -shift_size), slice(-shift_size, None)) w_slices = (slice(0, -window_size), slice(-window_size, -shift_size), slice(-shift_size, None)) cnt = 0 for h in h_slices: for w in w_slices: img_mask[:, h, w, :] = cnt cnt += 1 mask_windows = window_partition(img_mask, window_size) # nW, window_size, window_size, 1 mask_windows = mask_windows.view(-1, window_size * window_size) attn_mask = mask_windows.unsqueeze(1) - mask_windows.unsqueeze(2) attn_mask = attn_mask.masked_fill(attn_mask != 0, float(-100.0)).masked_fill(attn_mask == 0, float(0.0)) return attn_mask def cyclic_shift(x, shift_size): # 实现特征图的循环移位 shifted_x = torch.roll(x, shifts=(-shift_size, -shift_size), dims=(1, 2)) return shifted_x

4. 完整Swin Transformer Block实现

将W-MSA和SW-MSA组合成完整的Transformer Block：

class SwinTransformerBlock(nn.Module): def __init__(self, dim, num_heads, window_size=7, shift_size=0): super().__init__() self.dim = dim self.num_heads = num_heads self.window_size = window_size self.shift_size = shift_size self.norm1 = nn.LayerNorm(dim) self.attn = WindowAttention( dim, window_size=(self.window_size, self.window_size), num_heads=num_heads) self.norm2 = nn.LayerNorm(dim) self.mlp = nn.Sequential( nn.Linear(dim, dim * 4), nn.GELU(), nn.Linear(dim * 4, dim) ) if self.shift_size > 0: attn_mask = create_mask( window_size=self.window_size, shift_size=self.shift_size, H=56, W=56) # 示例中使用固定尺寸，实际应动态计算 else: attn_mask = None self.register_buffer("attn_mask", attn_mask) def forward(self, x): H, W = x.shape[1], x.shape[2] B, L, C = x.shape assert L == H * W, "输入特征尺寸不匹配" shortcut = x x = self.norm1(x) x = x.view(B, H, W, C) # 循环移位 if self.shift_size > 0: shifted_x = cyclic_shift(x, self.shift_size) else: shifted_x = x # 窗口划分 x_windows = window_partition(shifted_x, self.window_size) # nW*B, window_size, window_size, C x_windows = x_windows.view(-1, self.window_size * self.window_size, C) # nW*B, window_size*window_size, C # W-MSA/SW-MSA attn_windows = self.attn(x_windows, mask=self.attn_mask) # nW*B, window_size*window_size, C # 合并窗口 attn_windows = attn_windows.view(-1, self.window_size, self.window_size, C) shifted_x = window_reverse(attn_windows, self.window_size, H, W) # B H' W' C # 反向循环移位 if self.shift_size > 0: x = cyclic_shift(shifted_x, -self.shift_size) else: x = shifted_x x = x.view(B, H * W, C) # FFN x = shortcut + x x = x + self.mlp(self.norm2(x)) return x

关键设计要点：

1. 交替使用W-MSA和SW-MSA（通过shift_size控制）
2. 残差连接和LayerNorm的标准Transformer结构
3. 窗口划分与合并的逆操作
4. 循环移位的正反向处理

5. 实际应用中的性能优化技巧

在实现完整模型时，以下几个优化技巧可以显著提升效率：

内存优化技术：

高效注意力计算优化：

# 使用Flash Attention（如果可用） try: from flash_attn import flash_attention attn = flash_attention(q, k, v) except ImportError: # 回退到标准实现 attn = (q @ k.transpose(-2, -1)) * self.scale

分布式训练配置示例：

# 初始化分布式后端 torch.distributed.init_process_group(backend='nccl') # 包装模型 model = nn.parallel.DistributedDataParallel( model, device_ids=[local_rank], output_device=local_rank ) # 梯度聚合设置 optimizer = torch.optim.AdamW(model.parameters(), lr=1e-4) scaler = torch.cuda.amp.GradScaler()

在视觉任务实践中，Swin-Transformer的窗口注意力机制展现出几个显著优势：

1. 计算复杂度与图像大小呈线性关系而非平方关系
2. 通过层级设计适应多尺度特征提取
3. 移位窗口机制在保持效率的同时实现全局感受野
4. 相对位置编码更适合密集预测任务