别再死记硬背Swin Transformer结构了！用PyTorch手撕W-MSA和SW-MSA，彻底搞懂窗口注意力机制

从零实现Swin Transformer窗口注意力W-MSA与SW-MSA的PyTorch实战解析在计算机视觉领域Transformer架构正掀起一场革命。传统的卷积神经网络CNN长期主导着图像处理任务但Vision TransformerViT的出现打破了这一格局。然而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(fUsing 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_biasTrue, attn_drop0., proj_drop0.): 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, biasqkv_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, maskNone): 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×5612873.22×10⁹1.61×10⁹50%112×11212875.15×10¹⁰6.44×10⁹87.5%224×22412878.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_size7, shift_size0): 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_headsnum_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, maskattn_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(), cmapviridis) ax.set_title(fHead {i1}) ax.set_xticks([]) ax.set_yticks([]) fig.colorbar(im, axax) 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(), cmapgray) # 绘制水平线 for i in range(1, H // window_size): plt.axhline(yi * window_size - 0.5, colorred, linestyle-, linewidth2) # 绘制垂直线 for j in range(1, W // window_size): plt.axvline(xj * window_size - 0.5, colorred, linestyle-, linewidth2) plt.title(fWindow 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(), cmapviridis) plt.colorbar() plt.title(fRelative Position Index (Window Size: {window_size}x{window_size})) plt.show()