Swin Transformer实战：从零搭建移动窗口视觉模型（附PyTorch代码）

Swin Transformer实战：从零搭建移动窗口视觉模型（附PyTorch代码）

如果你在过去两年里关注过计算机视觉领域，大概率会反复听到一个词：Transformer。从自然语言处理领域“跨界”而来的Transformer架构，凭借其强大的全局建模能力，在图像分类任务上初露锋芒。然而，早期的视觉Transformer（如ViT）在处理高分辨率图像或密集预测任务（如目标检测、语义分割）时，往往会遇到两个棘手的问题：计算复杂度随图像尺寸呈二次方增长，以及缺乏像CNN那样的分层多尺度特征表示。这就像让一个擅长处理长篇文本的语言模型去逐像素分析一张4K高清图片，不仅效率低下，也难以捕捉从边缘到物体的多层次信息。

2021年，微软亚洲研究院提出的 Swin Transformer 巧妙地解决了这两个瓶颈。其核心创新在于 “移动窗口” 机制，将全局自注意力计算限制在局部非重叠窗口内，并通过在相邻层间滑动窗口来建立跨窗口连接，从而实现了线性计算复杂度和分层特征图。这一设计使其迅速成为视觉领域的通用骨干网络，在ImageNet分类、COCO目标检测、ADE20K语义分割等多个基准测试中刷新了记录。

对于开发者而言，理解Swin Transformer的原理是一回事，亲手将其实现并应用于自己的项目中则是另一回事。网络上充斥着大量理论解读，但关于如何从零开始用PyTorch搭建一个可运行的Swin Transformer模型，以及如何将其适配到自定义数据集的实战指南却相对零散。本文将扮演你的“代码搭档”，我们不仅会深入解析移动窗口注意力（Shifted Window based Self-Attention, SW-MSA）的高效批处理实现这一关键难点，还会提供完整的、模块化的PyTorch代码，并分享数据预处理、模型微调以及调试过程中可能遇到的“坑”及其解决方案。我们的目标很明确：让你在理解原理的基础上，获得立刻能跑起来的实践能力。

1. 环境准备与核心概念再梳理

在动手写代码之前，我们需要确保环境就绪，并对Swin Transformer的几个核心设计思想达成共识，这能帮助我们在实现时做出正确的设计决策。

开发环境配置 建议使用Python 3.8+和PyTorch 1.9+。我们将主要依赖以下库：

pip install torch torchvision
pip install timm  # 用于参考一些工具函数，如DropPath
pip install einops # 用于张量操作，让代码更清晰

使用Anaconda管理环境也是很好的选择。确保你的CUDA版本与PyTorch匹配，以启用GPU加速。

Swin Transformer的核心设计思想 与ViT将图像简单切割为16x16的图块（Patch）不同，Swin Transformer采用了更精细的层次化构建方式：

1. Patch Partition与Linear Embedding：输入图像（H x W x 3）首先被分割成4x4的非重叠图块（Patch），每个图块展平后的维度是443=48。随后通过一个线性层（Linear Embedding）将每个图块投影到一个预设的通道维度C（例如96）。这构成了“Stage 1”的输入。
2. Swin Transformer Block与层级下降：模型包含4个阶段（Stage）。每个Stage由若干Swin Transformer Block堆叠而成，并在两个Stage之间通过一个“Patch Merging”层来降低分辨率、增加通道数，从而构建金字塔式的特征层次。这与CNN中通过池化或步长卷积进行下采样的思路一脉相承。
3. 移动窗口自注意力（SW-MSA）：这是Swin的灵魂。在偶数层使用常规窗口划分（Window Multi-head Self-Attention, W-MSA），在奇数层使用偏移后的窗口划分（Shifted-Window MSA）。这种设计在保持窗口内计算高效性的同时，引入了相邻窗口间的信息交互。

提示：理解“移动窗口”的关键在于，它是对特征图进行划分，而不是对原始图像。经过Patch Merging后，特征图尺寸变小，但每个窗口所覆盖的原始图像感受野却在增大。

为了更直观地对比Swin Transformer与典型CNN（如ResNet）和早期ViT的结构差异，我们可以看下面这个表格：

特性	ResNet (CNN代表)	Vision Transformer (ViT)	Swin Transformer
核心操作	卷积、池化	全局自注意力	窗口自注意力 + 移动窗口
特征图	分层，多尺度	单一尺度（序列）	分层，多尺度
计算复杂度	线性（相对于图像尺寸）	二次方（相对于图块数）	线性（窗口大小固定）
归纳偏置	平移不变性、局部性	结构感知弱，需大量数据	局部性、层次性
适合任务	分类、检测、分割等	主要图像分类	分类、检测、分割等通用骨干
高分辨率适配	友好	困难（计算爆炸）	友好（线性复杂度）

这个对比清晰地展示了Swin Transformer如何取二者之长：既继承了CNN的多尺度层次化优点，又发挥了Transformer强大的建模能力，同时通过移动窗口机制规避了计算瓶颈。

2. 基础模块实现：从Patch Merging到窗口注意力

让我们从最基础的构件开始搭建。我们将采用自底向上的方式，先实现辅助模块，再组合成核心的Swin Transformer Block。

Patch Merging层 这个层位于两个Stage之间，功能类似于CNN中的池化层，用于下采样并增加通道数。其操作可以分解为：

• 提取局部块：将输入特征图在空间维度上，按2x2的邻域进行划分。
• 拼接与降维：将每个2x2区域内的4个特征向量在通道维度上拼接起来，这样通道数变为原来的4倍。
• 线性投影：通过一个线性层，将通道数降为原来的2倍（例如，从4C投影到2C）。同时，这实现了空间尺寸减半（H/2, W/2），通道数翻倍。

import torch
import torch.nn as nn
import torch.nn.functional as F

class PatchMerging(nn.Module):
    """
    下采样层，在Stage之间使用。
    输入: (B, H, W, C)
    输出: (B, H/2, W/2, 2C)
    """
    def __init__(self, input_resolution, dim):
        super().__init__()
        self.input_resolution = input_resolution
        self.dim = dim
        self.reduction = nn.Linear(4 * dim, 2 * dim, bias=False)
        self.norm = nn.LayerNorm(4 * dim)

    def forward(self, x):
        H, W = self.input_resolution
        B, L, C = x.shape
        assert L == H * W, "input feature has wrong size"
        assert H % 2 == 0 and W % 2 == 0, f"x size ({H}*{W}) are not even."

        x = x.view(B, H, W, C)

        # 提取2x2局部块并拼接
        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, 4*C)
        x = x.view(B, -1, 4 * C)  # (B, H/2 * W/2, 4*C)

        x = self.norm(x)
        x = self.reduction(x)  # (B, H/2 * W/2, 2*C)
        return x

窗口划分与还原 这是实现窗口注意力的关键辅助函数。我们需要能将一个批次的特征图序列（B, L, C）划分成多个窗口，并在计算后还原。

def window_partition(x, window_size):
    """
    将特征图划分为不重叠的窗口。
    输入:
        x: (B, H, W, C)
        window_size (int): 窗口大小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 (int): 窗口大小M
        H, W (int): 特征图的高和宽
    输出:
        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 x

相对位置偏置 在窗口自注意力中，为每个注意力头引入一个可学习的相对位置偏置（Relative Position Bias）被证明非常有效。它比绝对位置编码更适合视觉任务。

class RelativePositionBias(nn.Module):
    """ 可学习的相对位置偏置，为每个注意力头提供。"""
    def __init__(self, window_size, num_heads):
        super().__init__()
        self.window_size = window_size
        # 相对位置索引的范围是 [-M+1, M-1]，共 2M-1 个值
        self.relative_position_bias_table = nn.Parameter(
            torch.zeros((2 * window_size - 1) * (2 * window_size - 1), num_heads)
        )
        # 生成相对位置索引
        coords_h = torch.arange(window_size)
        coords_w = torch.arange(window_size)
        coords = torch.stack(torch.meshgrid([coords_h, coords_w], indexing='ij'))  # (2, M, M)
        coords_flatten = torch.flatten(coords, 1)  # (2, M*M)
        relative_coords = coords_flatten[:, :, None] - coords_flatten[:, None, :]  # (2, M*M, M*M)
        relative_coords = relative_coords.permute(1, 2, 0).contiguous()  # (M*M, M*M, 2)
        # 将坐标偏移到非负范围
        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)  # (M*M, M*M)
        self.register_buffer("relative_position_index", relative_position_index)

    def forward(self):
        # 根据索引从表中取出偏置值
        relative_position_bias = self.relative_position_bias_table[self.relative_position_index.view(-1)].view(
            self.window_size * self.window_size, self.window_size * self.window_size, -1
        )
        return relative_position_bias.permute(2, 0, 1).contiguous()  # (nH, M*M, M*M)

3. 核心突破：移动窗口自注意力（SW-MSA）的PyTorch实现

这是整个模型中最精妙也最具挑战性的部分。移动窗口的核心问题在于：偏移后，窗口数量会增加，且会出现小于MxM的不规则窗口。一种朴素的解决方案是给这些小窗口填充（Padding）到MxM，但这会带来大量的冗余计算。Swin Transformer论文中提出了一种循环移位（Cyclic Shift）结合掩码（Mask） 的高效批处理方法。

高效移动窗口自注意力模块 我们来实现这个模块。其流程可以概括为：

1. 对输入特征图进行循环移位。
2. 对移位后的特征图进行常规窗口划分。
3. 在每个窗口内计算带掩码的自注意力。
4. 计算结果并还原窗口。
5. 进行反向循环移位，将特征图恢复原状。

class WindowAttention(nn.Module):
    """ 基于窗口的多头自注意力，支持移位窗口（SW-MSA）和常规窗口（W-MSA）。"""
    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 = RelativePositionBias(window_size, num_heads)

        # 线性投影层
        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)

        self.softmax = nn.Softmax(dim=-1)

    def forward(self, x, mask=None):
        """
        输入:
            x: 输入特征图 (num_windows*B, N, C)，其中 N = M*M
            mask: (0/-inf) 掩码矩阵，用于SW-MSA，形状为 (nW, N, N) 或 None
        输出:
            x: 经过注意力加权后的特征 (num_windows*B, N, C)
        """
        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, head_dim)

        attn = (q @ k.transpose(-2, -1)) * self.scale  # (B_, nH, N, N)

        # 加上相对位置偏置
        relative_position_bias = self.relative_position_bias()
        attn = attn + relative_position_bias.unsqueeze(0)

        if mask is not None:
            # mask形状: (nW, N, N)，需要扩展到每个批次和每个头
            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 x

接下来，我们需要实现一个函数来为移位窗口生成注意力掩码。这个掩码用于在循环移位后，确保自注意力只计算在同一个子窗口内的像素之间。

def create_mask(H, W, window_size, shift_size):
    """
    为移位窗口自注意力生成掩码。
    输入:
        H, W: 特征图的高和宽
        window_size: 窗口大小M
        shift_size: 移位大小，通常为 M//2
    输出:
        mask: (nW, M*M, M*M)，其中nW是移位后窗口的数量
    """
    Hp = int(np.ceil(H / window_size)) * window_size
    Wp = int(np.ceil(W / window_size)) * window_size
    img_mask = torch.zeros((1, Hp, Wp, 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, M, M, 1)
    mask_windows = mask_windows.view(-1, window_size * window_size)  # (nW, M*M)
    attn_mask = mask_windows.unsqueeze(1) - mask_windows.unsqueeze(2)  # (nW, M*M, M*M)
    attn_mask = attn_mask.masked_fill(attn_mask != 0, float(-100.0)).masked_fill(attn_mask == 0, float(0.0))
    return attn_mask

现在，我们可以将上述所有组件组装成Swin Transformer Block。一个Block包含两个连续的子模块：第一个使用W-MSA，第二个使用SW-MSA。

class SwinTransformerBlock(nn.Module):
    def __init__(self, dim, input_resolution, num_heads, window_size=7, shift_size=0,
                 mlp_ratio=4., qkv_bias=True, drop=0., attn_drop=0., drop_path=0.):
        super().__init__()
        self.dim = dim
        self.input_resolution = input_resolution
        self.num_heads = num_heads
        self.window_size = window_size
        self.shift_size = shift_size
        self.mlp_ratio = mlp_ratio

        # 确保shift_size在合理范围内
        if min(self.input_resolution) <= self.window_size:
            self.shift_size = 0
            self.window_size = min(self.input_resolution)
        assert 0 <= self.shift_size < self.window_size, "shift_size must in 0-window_size"

        # 层归一化
        self.norm1 = nn.LayerNorm(dim)
        self.attn = WindowAttention(
            dim, window_size=self.window_size, num_heads=num_heads,
            qkv_bias=qkv_bias, attn_drop=attn_drop, proj_drop=drop
        )

        # DropPath (Stochastic Depth)
        self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()

        self.norm2 = nn.LayerNorm(dim)
        mlp_hidden_dim = int(dim * mlp_ratio)
        self.mlp = Mlp(in_features=dim, hidden_features=mlp_hidden_dim, drop=drop)

        # 为SW-MSA生成注意力掩码
        if self.shift_size > 0:
            H, W = self.input_resolution
            attn_mask = create_mask(H, W, self.window_size, self.shift_size)
        else:
            attn_mask = None
        self.register_buffer("attn_mask", attn_mask)

    def forward(self, x):
        H, W = self.input_resolution
        B, L, C = x.shape
        assert L == H * W, "input feature has wrong size"

        shortcut = x
        x = self.norm1(x)
        x = x.view(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, M, M, C)
        x_windows = x_windows.view(-1, self.window_size * self.window_size, C)  # (nW*B, M*M, C)

        # 窗口注意力
        attn_windows = self.attn(x_windows, mask=self.attn_mask)  # (nW*B, M*M, 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.view(B, H * W, C)

        # 第一个残差连接
        x = shortcut + self.drop_path(x)

        # MLP部分
        x = x + self.drop_path(self.mlp(self.norm2(x)))
        return x

这里用到了一个Mlp模块和DropPath函数，它们是Transformer中的标准组件：

class Mlp(nn.Module):
    def __init__(self, in_features, hidden_features=None, out_features=None, drop=0.):
        super().__init__()
        out_features = out_features or in_features
        hidden_features = hidden_features or in_features
        self.fc1 = nn.Linear(in_features, hidden_features)
        self.act = nn.GELU()
        self.fc2 = nn.Linear(hidden_features, out_features)
        self.drop = nn.Dropout(drop)

    def forward(self, x):
        x = self.fc1(x)
        x = self.act(x)
        x = self.drop(x)
        x = self.fc2(x)
        x = self.drop(x)
        return x

def drop_path(x, drop_prob: float = 0., training: bool = False):
    """Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks)."""
    if drop_prob == 0. or not training:
        return x
    keep_prob = 1 - drop_prob
    shape = (x.shape[0],) + (1,) * (x.ndim - 1)  # work with diff dim tensors, not just 2D ConvNets
    random_tensor = keep_prob + torch.rand(shape, dtype=x.dtype, device=x.device)
    random_tensor.floor_()  # binarize
    output = x.div(keep_prob) * random_tensor
    return output

class DropPath(nn.Module):
    def __init__(self, drop_prob=None):
        super(DropPath, self).__init__()
        self.drop_prob = drop_prob

    def forward(self, x):
        return drop_path(x, self.drop_prob, self.training)

4. 组装完整模型与在自定义数据集上微调

有了基础模块，我们现在可以像搭积木一样构建完整的Swin Transformer模型，并探讨如何将其应用于你自己的任务。

构建完整的Swin Transformer 模型由以下几个部分组成：

1. Patch Embedding：将图像分割为图块并进行线性投影。
2. 多个Stage：每个Stage包含若干Swin Transformer Block和一个可选的Patch Merging层（最后一个Stage后没有）。
3. 分类头：全局平均池化后接全连接层。

class SwinTransformer(nn.Module):
    def __init__(self, img_size=224, patch_size=4, in_chans=3, num_classes=1000,
                 embed_dim=96, depths=[2, 2, 6, 2], num_heads=[3, 6, 12, 24],
                 window_size=7, mlp_ratio=4., qkv_bias=True, drop_rate=0., attn_drop_rate=0.,
                 drop_path_rate=0.1, norm_layer=nn.LayerNorm, ape=False, **kwargs):
        super().__init__()

        self.num_classes = num_classes
        self.num_layers = len(depths)
        self.embed_dim = embed_dim
        self.ape = ape  # Absolute Position Embedding，Swin中通常不使用

        # 分割图块与线性嵌入
        self.patch_embed = PatchEmbed(
            img_size=img_size, patch_size=patch_size, in_chans=in_chans, embed_dim=embed_dim,
            norm_layer=norm_layer)
        num_patches = self.patch_embed.num_patches
        patches_resolution = self.patch_embed.patches_resolution
        self.patches_resolution = patches_resolution

        # 绝对位置编码（可选，通常效果不佳）
        if self.ape:
            self.absolute_pos_embed = nn.Parameter(torch.zeros(1, num_patches, embed_dim))
            nn.init.trunc_normal_(self.absolute_pos_embed, std=.02)

        self.pos_drop = nn.Dropout(p=drop_rate)

        # 随机深度衰减规则
        dpr = [x.item() for x in torch.linspace(0, drop_path_rate, sum(depths))]

        # 构建分层阶段
        self.layers = nn.ModuleList()
        for i_layer in range(self.num_layers):
            layer = BasicLayer(
                dim=int(embed_dim * 2 ** i_layer),
                input_resolution=(patches_resolution[0] // (2 ** i_layer),
                                 patches_resolution[1] // (2 ** i_layer)),
                depth=depths[i_layer],
                num_heads=num_heads[i_layer],
                window_size=window_size,
                mlp_ratio=mlp_ratio,
                qkv_bias=qkv_bias,
                drop=drop_rate,
                attn_drop=attn_drop_rate,
                drop_path=dpr[sum(depths[:i_layer]):sum(depths[:i_layer + 1])],
                downsample=PatchMerging if (i_layer < self.num_layers - 1) else None
            )
            self.layers.append(layer)

        self.norm = norm_layer(int(embed_dim * 2 ** (self.num_layers - 1)))
        self.avgpool = nn.AdaptiveAvgPool1d(1)
        self.head = nn.Linear(int(embed_dim * 2 ** (self.num_layers - 1)), num_classes) if num_classes > 0 else nn.Identity()

        self.apply(self._init_weights)

    def _init_weights(self, m):
        if isinstance(m, nn.Linear):
            nn.init.trunc_normal_(m.weight, std=.02)
            if isinstance(m, nn.Linear) and m.bias is not None:
                nn.init.constant_(m.bias, 0)
        elif isinstance(m, nn.LayerNorm):
            nn.init.constant_(m.bias, 0)
            nn.init.constant_(m.weight, 1.0)

    def forward_features(self, x):
        x = self.patch_embed(x)
        if self.ape:
            x = x + self.absolute_pos_embed
        x = self.pos_drop(x)

        for layer in self.layers:
            x = layer(x)

        x = self.norm(x)  # (B, L, C)
        x = self.avgpool(x.transpose(1, 2))  # (B, C, 1)
        x = torch.flatten(x, 1)
        return x

    def forward(self, x):
        x = self.forward_features(x)
        x = self.head(x)
        return x

其中，PatchEmbed和BasicLayer是辅助构建块：

class PatchEmbed(nn.Module):
    def __init__(self, img_size=224, patch_size=4, in_chans=3, embed_dim=96, norm_layer=None):
        super().__init__()
        img_size = (img_size, img_size)
        patch_size = (patch_size, patch_size)
        patches_resolution = [img_size[0] // patch_size[0], img_size[1] // patch_size[1]]
        self.img_size = img_size
        self.patch_size = patch_size
        self.patches_resolution = patches_resolution
        self.num_patches = patches_resolution[0] * patches_resolution[1]

        self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=patch_size)
        if norm_layer is not None:
            self.norm = norm_layer(embed_dim)
        else:
            self.norm = None

    def forward(self, x):
        B, C, H, W = x.shape
        assert H == self.img_size[0] and W == self.img_size[1], \
            f"Input image size ({H}*{W}) doesn't match model ({self.img_size[0]}*{self.img_size[1]})."
        x = self.proj(x).flatten(2).transpose(1, 2)  # (B, Ph*Pw, C)
        if self.norm is not None:
            x = self.norm(x)
        return x

class BasicLayer(nn.Module):
    def __init__(self, dim, input_resolution, depth, num_heads, window_size,
                 mlp_ratio=4., qkv_bias=True, drop=0., attn_drop=0.,
                 drop_path=0., downsample=None):
        super().__init__()
        self.dim = dim
        self.input_resolution = input_resolution
        self.depth = depth
        self.blocks = nn.ModuleList([
            SwinTransformerBlock(
                dim=dim, input_resolution=input_resolution,
                num_heads=num_heads, window_size=window_size,
                shift_size=0 if (i % 2 == 0) else window_size // 2,  # 交替使用W-MSA和SW-MSA
                mlp_ratio=mlp_ratio, qkv_bias=qkv_bias,
                drop=drop, attn_drop=attn_drop,
                drop_path=drop_path[i] if isinstance(drop_path, list) else drop_path,
            )
            for i in range(depth)
        ])
        if downsample is not None:
            self.downsample = downsample(input_resolution, dim=dim)
        else:
            self.downsample = None

    def forward(self, x):
        for blk in self.blocks:
            x = blk(x)
        if self.downsample is not None:
            x = self.downsample(x)
        return x

至此，一个完整的Swin-Tiny（Swin-T）模型就构建完成了。你可以通过调整embed_dim, depths, num_heads等参数来定义Swin-Small, Swin-Base, Swin-Large等变体。

在自定义数据集上微调 假设你有一个自己的图像分类数据集，目录结构如下：

your_dataset/
├── train/
│   ├── class1/
│   │   ├── img1.jpg
│   │   └── ...
│   ├── class2/
│   └── ...
└── val/
    ├── class1/
    ├── class2/
    └── ...

微调流程通常包括以下步骤：

1. 数据加载与增强：使用torchvision.datasets.ImageFolder和torchvision.transforms。
2. 加载预训练权重：从官方仓库或timm库加载在ImageNet-1K或ImageNet-22K上预训练的权重。
3. 修改分类头：将最后的全连接层输出维度改为你的类别数。
4. 选择优化器与学习率调度：通常使用AdamW优化器，并采用带热身的余弦退火学习率。
5. 训练与验证：在训练集上微调，在验证集上监控指标。

下面是一个简化的微调代码框架：

import torch
from torch import nn, optim
from torch.utils.data import DataLoader
from torchvision import datasets, transforms
from timm.data import create_transform
from timm.models import swin_tiny_patch4_window7_224
import os

# 1. 数据预处理
train_transform = create_transform(
    input_size=224,
    is_training=True,
    color_jitter=0.4,
    auto_augment='rand-m9-mstd0.5-inc1',
    interpolation='bicubic',
    re_prob=0.25,
    re_mode='pixel',
    re_count=1,
)
val_transform = transforms.Compose([
    transforms.Resize(256),
    transforms.CenterCrop(224),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])

train_dataset = datasets.ImageFolder(root='your_dataset/train', transform=train_transform)
val_dataset = datasets.ImageFolder(root='your_dataset/val', transform=val_transform)

train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True, num_workers=4, pin_memory=True)
val_loader = DataLoader(val_dataset, batch_size=32, shuffle=False, num_workers=4, pin_memory=True)

# 2. 加载模型并修改分类头
model = swin_tiny_patch4_window7_224(pretrained=True)  # 使用timm加载预训练模型
num_ftrs = model.head.in_features
model.head = nn.Linear(num_ftrs, len(train_dataset.classes))  # 修改分类头

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = model.to(device)

# 3. 定义损失函数和优化器
criterion = nn.CrossEntropyLoss()
optimizer = optim.AdamW(model.parameters(), lr=1e-4, weight_decay=0.05)

# 4. 训练循环 (简化版)
num_epochs = 30
for epoch in range(num_epochs):
    model.train()
    running_loss = 0.0
    for inputs, labels in train_loader:
        inputs, labels = inputs.to(device), labels.to(device)
        optimizer.zero_grad()
        outputs = model(inputs)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()
        running_loss += loss.item()
    print(f'Epoch {epoch+1}, Loss: {running_loss/len(train_loader):.4f}')

    # 验证
    model.eval()
    correct = 0
    total = 0
    with torch.no_grad():
        for inputs, labels in val_loader:
            inputs, labels = inputs.to(device), labels.to(device)
            outputs = model(inputs)
            _, predicted = torch.max(outputs.data, 1)
            total += labels.size(0)
            correct += (predicted == labels).sum().item()
    print(f'Validation Accuracy: {100 * correct / total:.2f}%')

注意：在实际项目中，你需要更完善的训练循环，包括学习率调度（如CosineAnnealingLR）、模型保存、日志记录、混合精度训练（torch.cuda.amp）以及更复杂的数据增强策略。对于目标检测或分割任务，你需要将Swin Transformer作为骨干网络集成到如Mask R-CNN、UperNet等框架中，这通常涉及修改特征金字塔网络（FPN）的输入。

5. 实战技巧、常见问题与性能优化指南

将模型跑起来只是第一步，要让其在你的任务上达到最佳性能，并高效地部署，还需要掌握一些实战技巧。

数据预处理与增强策略 Swin Transformer对数据增强比较敏感。除了标准的随机裁剪、水平翻转外，论文中使用了如RandAugment、MixUp、CutMix等强增强方法。对于自定义数据集，尤其是数据量不大时，这些增强能有效防止过拟合。你可以使用timm库中的create_transform来方便地配置，或者使用albumentations库进行更灵活的增强。

学习率与优化器设置

• 优化器：AdamW 是训练Transformer类模型的首选，其解耦了权重衰减，效果通常优于SGD。
• 学习率：微调时，学习率应设置得比从头训练小。一个常见的策略是：lr = base_lr * (batch_size / 256)。对于Swin-T微调，base_lr在1e-4到5e-4之间是个不错的起点。
• 学习率调度：带热身的余弦退火（Cosine Annealing with Warmup）非常有效。热身阶段（例如5个epoch）让学习率从一个小值线性增长到初始学习率，有助于训练稳定。
• 权重衰减：通常设置为0.05。
• 梯度裁剪：可以设置一个最大值（如max_norm=1.0）来防止梯度爆炸，这在训练深度Transformer时有时是必要的。

混合精度训练（AMP） 使用Automatic Mixed Precision可以显著减少显存占用并加快训练速度，几乎不影响精度。

from torch.cuda.amp import autocast, GradScaler

scaler = GradScaler()
for inputs, labels in train_loader:
    optimizer.zero_grad()
    with autocast():
        outputs = model(inputs)
        loss = criterion(outputs, labels)
    scaler.scale(loss).backward()
    scaler.step(optimizer)
    scaler.update()

模型保存与加载 保存最佳验证集性能的模型，并保存优化器状态以便恢复训练。

if val_acc > best_acc:
    best_acc = val_acc
    torch.save({
        'epoch': epoch,
        'model_state_dict': model.state_dict(),
        'optimizer_state_dict': optimizer.state_dict(),
        'best_acc': best_acc,
    }, 'best_swin_model.pth')

常见报错与解决方案

1. CUDA out of memory：这是最常见的问题。

   • 降低批次大小（Batch Size）：这是最直接的方法。
   • 使用梯度累积：如果硬件限制导致批次大小只能设为1或2，可以通过梯度累积来模拟更大的批次。例如，每4个批次才更新一次权重（accumulation_steps=4），等效于批次大小扩大4倍。
   • 使用混合精度训练（AMP）：如前所述，可以节省显存。
   • 检查输入图像尺寸：确保输入尺寸与模型预期一致（默认224x224）。更大的尺寸会显著增加显存消耗。
   • 使用torch.utils.checkpoint：即激活检查点，用计算时间换显存空间，适用于非常深的模型。

2. 训练损失不下降或准确率极低：

   • 检查数据加载：确认数据路径、标签是否正确，可视化一些样本看看增强是否合理。
   • 检查学习率：学习率可能太大（导致震荡）或太小（导致收敛慢）。尝试使用学习率查找器（如torch-lr-finder）找到一个合适的范围。
   • 检查预训练权重加载：确认是否成功加载了预训练权重。打印模型参数，看看除了分类头以外的层是否有非随机初始化的值。
   • 冻结骨干网络：在微调初期，可以尝试先冻结Swin Transformer的主干部分，只训练分类头几个epoch，让模型先适应新的数据分布，然后再解冻全部参数进行微调。

3. 验证集性能远差于训练集（过拟合）：

   • 加强数据增强：使用更激进的数据增强，如RandAugment、CutMix、MixUp。
   • 增加正则化：增大Dropout率、DropPath率，或使用更强的权重衰减。
   • 获取更多数据：如果可能，收集更多训练样本。
   • 早停（Early Stopping）：监控验证集损失，当其在连续多个epoch不再下降时停止训练。

性能优化与推理加速

• TorchScript或ONNX导出：将训练好的PyTorch模型导出为TorchScript或ONNX格式，可以利用PyTorch的JIT编译器或专门的推理引擎（如ONNX Runtime, TensorRT）进行优化，提升推理速度。
• 使用半精度（FP16）推理：在推理时使用model.half()并将输入数据转换为半精度，可以加速计算并减少显存占用，对精度影响通常很小。
• 针对移动端部署：可以考虑使用模型压缩技术，如知识蒸馏（用大模型Swin-L教一个小模型）、剪枝（移除不重要的权重）或量化（将FP32权重转换为INT8），以大幅减少模型大小和计算量。PyTorch提供了torch.quantization工具进行量化。

一个实用的调试清单 在开始长时间训练前，快速运行以下检查：

• [ ] 输入数据维度是否正确？(B, C, H, W)
• [ ] 标签是否在有效范围内？(0 到 num_classes-1)
• [ ] 前向传播能否在不使用GPU的情况下跑通一个小批次？
• [ ] 损失函数计算是否正常？
• [ ] 反向传播是否能成功计算梯度？
• [ ] 优化器步骤是否能更新参数？
• [ ] 在单个批次上，训练模式（model.train()）和评估模式（model.eval()）的输出是否有差异（Dropout/BatchNorm导致）？

从理解移动窗口的巧妙设计，到亲手实现其高效批处理逻辑，再到将其应用于实际任务并解决可能遇到的各种工程问题，这个过程本身就是一个深度学习工程师能力成长的缩影。Swin Transformer的成功不仅仅在于其出色的性能，更在于它为我们提供了一个如何将Transformer高效适配视觉任务的经典范式。希望这份结合了原理、代码和实战经验的指南，能成为你探索视觉Transformer世界的一块坚实垫脚石。在实际项目中，多实验、多分析、多总结，才是掌握任何新架构的不二法门。