从ViT到Swin：手把手教你用PyTorch复现Swin Transformer的窗口注意力机制-开发者社区

从ViT到Swin：手把手教你用PyTorch复现Swin Transformer的窗口注意力机制

在计算机视觉领域，Transformer架构正逐渐取代传统的CNN成为新的主流。然而，标准的Vision Transformer（ViT）在处理高分辨率图像时面临着计算复杂度平方级增长的挑战。Swin Transformer通过引入**窗口注意力（W-MSA）和移位窗口注意力（SW-MSA）**机制，不仅显著降低了计算量，还构建了层次化的特征表示。本文将深入剖析这两种注意力机制的原理与实现细节，带你从零开始用PyTorch实现这一创新设计。

1. 窗口注意力机制的核心思想

传统ViT的全局自注意力计算复杂度与图像分辨率呈平方关系，这使得它难以处理大尺寸图像。Swin Transformer的突破在于将图像划分为不重叠的局部窗口，在每个窗口内独立计算自注意力。这种设计带来了两个关键优势：

计算复杂度线性增长：假设图像被划分为M×M的窗口，计算复杂度从O(H²W²)降至O(M²HW)，其中H和W是特征图的高和宽
局部归纳偏置：类似于CNN的局部感受野，更符合视觉任务的特性

窗口划分的实现可以通过window_partition函数完成：

def window_partition(x, window_size): 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

注意：在实际应用中，窗口大小通常设置为7×7，这是一个在计算效率和模型性能之间取得良好平衡的经验值

2. W-MSA的PyTorch实现详解

窗口多头自注意力（W-MSA）是Swin Transformer的基础模块，其实现需要考虑以下几个关键点：

相对位置编码：由于窗口内的像素位置关系比绝对位置更重要，Swin采用了可学习的相对位置偏置
注意力掩码：在移位窗口注意力中需要特殊的掩码机制
多头注意力计算：标准的缩放点积注意力计算

以下是完整的WindowAttention模块实现：

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) 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_, nH, N, C q = q * self.scale attn = (q @ k.transpose(-2, -1)) # B_, nH, 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) 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

这段代码中有几个关键点值得特别关注：

相对位置编码：通过relative_position_bias_table实现，为每个相对位置组合学习一个偏置项
高效索引计算：使用巧妙的坐标变换生成位置索引，避免重复计算
掩码支持：预留了mask参数接口，为后续的移位窗口注意力做准备

3. 移位窗口注意力（SW-MSA）的实现技巧

单纯的窗口注意力虽然降低了计算复杂度，但也限制了跨窗口的信息交互。Swin Transformer通过移位窗口划分的创新设计解决了这一问题。具体实现分为三个步骤：

常规窗口划分：与W-MSA相同的方式处理输入特征
窗口移位：将特征图在水平和垂直方向各移位(window_size//2)个像素
掩码处理：确保注意力只在正确的窗口内计算

移位窗口的实现需要特殊的掩码机制，以下是生成掩码的关键代码：

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

移位窗口注意力的前向传播过程如下：

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, window_size=(self.window_size, self.window_size), num_heads=num_heads) if self.shift_size > 0: # 只在SW-MSA时计算掩码 H, W = 56, 56 # 假设输入特征图大小 img_mask = torch.zeros((1, H, W, 1)) h_slices = (slice(0, -self.window_size), slice(-self.window_size, -self.shift_size), slice(-self.shift_size, None)) w_slices = (slice(0, -self.window_size), slice(-self.window_size, -self.shift_size), slice(-self.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, self.window_size) mask_windows = mask_windows.view(-1, self.window_size * self.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)) else: attn_mask = None self.register_buffer("attn_mask", attn_mask) def forward(self, x): H, W = x.shape[2], x.shape[3] x = x.permute(0, 2, 3, 1) # B H W C # 循环移位 if self.shift_size > 0: shifted_x = torch.roll(x, shifts=(-self.shift_size, -self.shift_size), dims=(1, 2)) 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 = torch.roll(shifted_x, shifts=(self.shift_size, self.shift_size), dims=(1, 2)) else: x = shifted_x x = x.permute(0, 3, 1, 2) # B C H W return x

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

在实际项目中应用Swin Transformer的窗口注意力机制时，以下几个优化技巧可以显著提升模型效率：

内存优化策略：

优化方法	实现方式	预期收益
梯度检查点	使用torch.utils.checkpoint	减少显存占用30-50%
混合精度训练	启用AMP自动混合精度	加速训练20-30%
内存高效注意力	实现内存优化的注意力计算	降低峰值显存需求

计算加速技巧：

# 启用PyTorch的优化设置 torch.backends.cuda.matmul.allow_tf32 = True # 启用TF32加速 torch.backends.cudnn.benchmark = True # 启用cuDNN自动调优 # 混合精度训练示例 scaler = torch.cuda.amp.GradScaler() with torch.cuda.amp.autocast(): output = model(input) loss = criterion(output, target) scaler.scale(loss).backward() scaler.step(optimizer) scaler.update()

自定义数据集适配建议：