从零实现Swin Transformer窗口注意力：W-MSA与SW-MSA的PyTorch实战解析
在计算机视觉领域，Transformer架构正掀起一场革命。传统的卷积神经网络（CNN）长期主导着图像处理任务，但Vision Transformer（ViT）的出现打破了这一格局。然而，ViT在处理高分辨率图像时面临计算复杂度平方级增长的挑战。Swin Transformer通过引入**窗口多头自注意力（W-MSA）和移位窗口多头自注意力（SW-MSA）**机制，巧妙地解决了这一问题，成为视觉Transformer发展的重要里程碑。
本文将带您深入Swin Transformer的核心机制，通过PyTorch代码实现W-MSA和SW-MSA模块，彻底理解窗口注意力的工作原理。不同于单纯的理论讲解，我们将从代码层面拆解每个关键步骤，包括窗口划分、相对位置偏置、掩码计算等核心环节，让抽象的概念变得具体可操作。
1. 环境准备与基础模块搭建
在开始实现W-MSA和SW-MSA之前，我们需要搭建好开发环境并准备一些基础工具函数。以下是推荐的开发环境配置：
import torch import torch.nn as nn import torch.nn.functional as F import matplotlib.pyplot as plt import numpy as np # 检查GPU可用性 device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') print(f"Using device: {device}")窗口划分是W-MSA的基础操作，我们先实现一个通用的窗口划分函数：
def window_partition(x, window_size): """ 将输入特征图划分为不重叠的窗口 参数: x: (B, H, W, C) window_size: 窗口大小(M) 返回: 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对应的窗口还原函数：
def window_reverse(windows, window_size, H, W): """ 将划分的窗口还原为原始特征图 参数: windows: (num_windows*B, window_size, window_size, C) window_size: 窗口大小(M) H: 特征图高度 W: 特征图宽度 返回: x: (B, H, W, C) """ B = int(windows.shape[0] / (H * W / window_size / window_size)) x = windows.view(B, H // window_size, W // window_size, window_size, window_size, -1) x = x.permute(0, 1, 3, 2, 4, 5).contiguous().view(B, H, W, -1) return x2. W-MSA的实现与优化
窗口多头自注意力（W-MSA）是Swin Transformer的核心创新之一，它通过将特征图划分为不重叠的窗口，在每个窗口内独立计算自注意力，大幅降低了计算复杂度。
2.1 基础W-MSA实现
我们先实现一个基础的窗口注意力模块：
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 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, bias=qkv_bias) self.attn_drop = nn.Dropout(attn_drop) self.proj = nn.Linear(dim, dim) self.proj_drop = nn.Dropout(proj_drop) nn.init.trunc_normal_(self.relative_position_bias_table, std=.02) self.softmax = nn.Softmax(dim=-1)2.2 前向传播实现
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) q = q * self.scale attn = (q @ k.transpose(-2, -1)) # (B_, num_heads, N, N) 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 = self.softmax(attn) else: attn = self.softmax(attn) 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 x2.3 计算复杂度分析
W-MSA的计算量优势主要体现在处理大尺寸特征图时。让我们通过具体数据对比W-MSA和普通MSA的计算量：
| 特征图尺寸 | 通道数 | 窗口大小 | MSA计算量 (FLOPs) | W-MSA计算量 (FLOPs) | 节省比例 |
|---|---|---|---|---|---|
| 56×56 | 128 | 7 | 3.22×10⁹ | 1.61×10⁹ | 50% |
| 112×112 | 128 | 7 | 5.15×10¹⁰ | 6.44×10⁹ | 87.5% |
| 224×224 | 128 | 7 | 8.24×10¹¹ | 1.03×10¹¹ | 87.5% |
从表中可以看出，随着特征图尺寸增大，W-MSA节省的计算量呈平方级增长。
3. SW-MSA的实现与高效计算
虽然W-MSA降低了计算复杂度，但它限制了窗口间的信息交流。移位窗口多头自注意力（SW-MSA）通过周期性移动窗口位置，实现了跨窗口连接，同时保持了计算效率。
3.1 窗口移位实现
def create_mask(H, W, window_size, shift_size): # 创建用于SW-MSA的掩码 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_mask3.2 高效批处理实现
SW-MSA的关键挑战在于移位后窗口数量增加和形状不规则。我们通过循环移位和掩码技术实现高效计算：
def cyclic_shift(x, shift_size): # 实现特征图的循环移位 if shift_size > 0: shifted_x = torch.roll(x, shifts=(-shift_size, -shift_size), dims=(1, 2)) else: shifted_x = x return shifted_x def reverse_cyclic_shift(shifted_x, shift_size): # 反向循环移位，恢复原始布局 if shift_size > 0: x = torch.roll(shifted_x, shifts=(shift_size, shift_size), dims=(1, 2)) else: x = shifted_x return x3.3 SW-MSA完整流程
结合上述组件，SW-MSA的完整实现流程如下：
- 对输入特征图进行循环移位
- 在移位后的特征图上划分窗口
- 计算带掩码的窗口注意力
- 合并窗口并反向循环移位
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.attn = WindowAttention( dim, window_size=(self.window_size, self.window_size), num_heads=num_heads ) # 初始化MLP层等其它组件... def forward(self, x): H, W = x.shape[1], x.shape[2] B, L, C = x.shape # 循环移位 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 # 创建注意力掩码 if self.shift_size > 0: attn_mask = create_mask(H, W, self.window_size, self.shift_size) else: attn_mask = None # 窗口注意力计算 attn_windows = self.attn(x_windows, mask=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 = reverse_cyclic_shift(shifted_x, self.shift_size) else: x = shifted_x return x4. 可视化分析与调试技巧
为了深入理解W-MSA和SW-MSA的工作原理，可视化分析是极其有效的手段。下面介绍几种实用的可视化方法。
4.1 注意力权重可视化
def visualize_attention(attention_weights, window_size): """ 可视化注意力权重 参数: attention_weights: (num_heads, window_size*window_size, window_size*window_size) window_size: 窗口大小 """ num_heads = attention_weights.shape[0] fig, axes = plt.subplots(1, num_heads, figsize=(20, 5)) if num_heads == 1: axes = [axes] for i in range(num_heads): ax = axes[i] im = ax.imshow(attention_weights[i].detach().cpu().numpy(), cmap='viridis') ax.set_title(f'Head {i+1}') ax.set_xticks([]) ax.set_yticks([]) fig.colorbar(im, ax=ax) plt.tight_layout() plt.show()4.2 窗口划分可视化
def visualize_window_partition(feature_map, window_size): """ 可视化特征图的窗口划分 参数: feature_map: (H, W, C) window_size: 窗口大小 """ H, W, _ = feature_map.shape plt.figure(figsize=(10, 10)) plt.imshow(feature_map.mean(dim=-1).detach().cpu().numpy(), cmap='gray') # 绘制水平线 for i in range(1, H // window_size): plt.axhline(y=i * window_size - 0.5, color='red', linestyle='-', linewidth=2) # 绘制垂直线 for j in range(1, W // window_size): plt.axvline(x=j * window_size - 0.5, color='red', linestyle='-', linewidth=2) plt.title(f'Window Partition (Window Size: {window_size}x{window_size})') plt.axis('off') plt.show()4.3 调试技巧与常见问题
在实现W-MSA和SW-MSA时，经常会遇到以下问题及解决方案：
形状不匹配错误：
- 检查窗口划分和还原过程中张量的形状变化
- 确保注意力权重的形状为
(num_windows*B, num_heads, window_size*window_size, window_size*window_size)
梯度消失或爆炸：
- 初始化相对位置偏置表时使用较小的标准差（如0.02）
- 在注意力计算后添加LayerNorm
掩码不起作用：
- 确保掩码值足够大（如-100），使得softmax后对应位置的权重接近0
- 检查掩码是否正确地应用于注意力权重
性能问题：
- 对于大尺寸输入，考虑使用混合精度训练
- 在SW-MSA中，合理选择移位大小（通常为窗口大小的一半）
# 调试示例：检查相对位置索引 def check_relative_position_index(window_size): coords_h = torch.arange(window_size) coords_w = torch.arange(window_size) coords = torch.stack(torch.meshgrid([coords_h, coords_w])) 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 relative_position_index = relative_coords.sum(-1) print("相对位置索引矩阵：") print(relative_position_index) # 可视化 plt.figure(figsize=(8, 6)) plt.imshow(relative_position_index.numpy(), cmap='viridis') plt.colorbar() plt.title(f'Relative Position Index (Window Size: {window_size}x{window_size})') plt.show()
计算)
