从Transformer到ViT与Swin详解

从Transformer到ViT与Swin：视觉Transformer的演进与实战

• 前言
• 一、Transformer
• • [1.1 整体架构](#1.1 整体架构)
  • [1.2 核心组件详解](#1.2 核心组件详解)
  • • [(1) 自注意力机制（Self-Attention）](#(1) 自注意力机制（Self-Attention）)
    • [(2) 多头注意力（Multi-Head Attention）](#(2) 多头注意力（Multi-Head Attention）)
    • [(3) 位置编码（Positional Encoding）](#(3) 位置编码（Positional Encoding）)
    • [(4) 前馈网络（Feed Forward Network）](#(4) 前馈网络（Feed Forward Network）)
    • [(5) 残差连接与层归一化](#(5) 残差连接与层归一化)
  • [1.3 训练与推理要点](#1.3 训练与推理要点)
• [二、Vision Transformer (ViT)](#二、Vision Transformer (ViT))
• • [2.1 设计动机](#2.1 设计动机)
  • [2.2 模型结构](#2.2 模型结构)
  • [2.3 关键设计细节](#2.3 关键设计细节)
  • [2.4 训练策略](#2.4 训练策略)
  • [2.5 优缺点](#2.5 优缺点)
• [三、Swin Transformer](#三、Swin Transformer)
• • [3.1 设计动机](#3.1 设计动机)
  • [3.2 核心创新](#3.2 核心创新)
  • • [(1) 层级化架构（Hierarchical Feature Maps）](#(1) 层级化架构（Hierarchical Feature Maps）)
    • [(2) 窗口多头自注意力（W-MSA）](#(2) 窗口多头自注意力（W-MSA）)
    • [(3) 移位窗口注意力（SW-MSA）](#(3) 移位窗口注意力（SW-MSA）)
  • [3.3 相对位置编码](#3.3 相对位置编码)
  • [3.4 整体配置（Swin-Tiny为例）](#3.4 整体配置（Swin-Tiny为例）)
• 四、模型对比
• 五、代码实战
• • [5.1 环境准备](#5.1 环境准备)
  • [5.2 Transformer核心模块实现](#5.2 Transformer核心模块实现)
  • [5.3 ViT完整实现](#5.3 ViT完整实现)
  • [5.4 Swin Transformer核心组件实现](#5.4 Swin Transformer核心组件实现)
  • [5.5 训练与推理代码](#5.5 训练与推理代码)
• 六、实践建议与常见问题
• • [6.1 训练技巧](#6.1 训练技巧)
  • [6.2 常见问题及解决](#6.2 常见问题及解决)
  • [6.3 模型选择指南](#6.3 模型选择指南)

前言

2017年，Attention Is All You Need一文提出了Transformer架构，彻底改变了自然语言处理领域。2020年，Vision Transformer（ViT）将纯Transformer结构直接应用于图像分类，展示了Transformer在视觉任务中的巨大潜力。2021年，Swin Transformer引入了层级化设计和窗口注意力机制，使Transformer在多种视觉任务上全面超越CNN。

本文将从模型结构、原理到代码实现，系统介绍这三个关键模型，帮助读者建立起从NLP到CV的Transformer知识体系。

---

一、Transformer

1.1 整体架构

Transformer完全基于注意力机制，抛弃了RNN的循环结构和CNN的卷积结构。其架构由编码器（Encoder）和解码器（Decoder）两部分组成，每部分由多个相同的层堆叠而成。

原始Transformer架构：

• 编码器：N=6层，每层包含 Multi-Head Self-Attention + Feed Forward Network
• 解码器：N=6层，每层包含 Masked Self-Attention + Cross-Attention + FFN

1.2 核心组件详解

(1) 自注意力机制（Self-Attention）

自注意力是Transformer的核心。对于输入序列中的每个位置，它计算该位置与所有位置的相关性，然后加权求和得到新的表示。

计算公式：

Attention(Q, K, V) = softmax(Q × K^T / √d_k) × V

其中Q（查询）、K（键）、V（值）均由输入通过线性变换得到。除以√d_k是为了防止点积过大导致softmax梯度消失。

理解要点：

• Q和K的点积衡量两个位置之间的"相关性"或"注意力分数"
• softmax将分数归一化为概率分布
• 最终输出是V的加权和，权重来自注意力分数

(2) 多头注意力（Multi-Head Attention）

将输入投影到h个不同的子空间，分别执行注意力，最后拼接并投影。

MultiHead(Q,K,V) = Concat(head_1, ..., head_h) × W_O
head_i = Attention(Q × W_i^Q, K × W_i^K, V × W_i^V)

多头机制允许模型在不同子空间关注不同位置，增强了表达能力。

(3) 位置编码（Positional Encoding）

由于自注意力本身不具备顺序信息，需要显式注入位置信息。原始Transformer使用正弦/余弦函数：

PE(pos, 2i) = sin(pos / 10000^(2i/d_model))
PE(pos, 2i+1) = cos(pos / 10000^(2i/d_model))

(4) 前馈网络（Feed Forward Network）

每个注意力层后接一个两层的全连接网络：

FFN(x) = max(0, x × W_1 + b_1) × W_2 + b_2

(5) 残差连接与层归一化

每个子层都使用残差连接和层归一化：

output = LayerNorm(x + Sublayer(x))

1.3 训练与推理要点

• 训练：使用Teacher Forcing，解码器输入为真实目标序列（右移一位），采用masked attention防止看到未来信息
• 推理：自回归生成，每次输出一个token，作为下一步输入

---

二、Vision Transformer (ViT)

2.1 设计动机

ViT的核心理念是：将图像视为一系列patch的序列，直接用Transformer处理。这一想法颠覆了"图像必须用CNN"的传统认知。

2.2 模型结构

ViT结构（以ViT-B/16为例）：

1. 输入图像: 224×224×3
2. 分块: 16×16的patch，得到14×14=196个patch
3. 线性投影: 每个patch展平为16×16×3=768维，通过线性层映射到D维
4. 添加位置编码: 可学习的位置编码
5. 添加CLS token: 用于分类的额外可学习token
6. 经过L层Transformer Encoder（L=12 for ViT-B）
7. 分类头: 取CLS token输出，通过MLP得到类别

2.3 关键设计细节

Patch Embedding实现：

# 等价于卷积：kernel_size=patch_size, stride=patch_size
patch_emb = nn.Conv2d(in_chans, embed_dim, 
                      kernel_size=patch_size, stride=patch_size)

Position Embedding：使用可学习的1D位置编码，而非正弦编码。因为patch的顺序是规则的1D序列。

CLS Token：借鉴BERT，在序列开头添加一个可学习的token，其最终输出用于分类。也可以使用所有token的平均池化，但CLS已被广泛采用。

2.4 训练策略

ViT需要大规模预训练（如JFT-300M或ImageNet-21K），在小数据集上表现不如CNN。关键原因是Transformer缺乏CNN的归纳偏置（局部性、平移等变性），需要更多数据学习这些模式。

推荐训练配置：

• 优化器：AdamW
• 学习率：warmup + cosine decay
• 正则化：weight decay, dropout, label smoothing
• 数据增强：RandAugment, MixUp, CutMix

2.5 优缺点

优点	缺点
架构简洁，无卷积归纳偏置	需要大量训练数据
全局感受野，擅长建模长程依赖	计算复杂度与patch数平方相关
易于扩展，参数高效	对分辨率变化敏感

---

三、Swin Transformer

3.1 设计动机

ViT虽然在分类任务上表现优异，但作为通用视觉骨干存在两个问题：

1. 计算复杂度：与图像尺寸平方相关，难以处理高分辨率
2. 缺乏多尺度特征：不适合检测、分割等需要多尺度信息的任务

Swin Transformer通过层级化设计 和窗口注意力解决了上述问题。

3.2 核心创新

(1) 层级化架构（Hierarchical Feature Maps）

Swin通过patch合并（Patch Merging）逐步降低空间分辨率、增加通道数，形成类似CNN的特征金字塔：

输入: H×W×3
Stage 1: Patch Partition → H/4 × W/4 × C (Linear Embedding)
Stage 2: Patch Merging → H/8 × W/8 × 2C
Stage 3: Patch Merging → H/16 × W/16 × 4C
Stage 4: Patch Merging → H/32 × W/32 × 8C

(2) 窗口多头自注意力（W-MSA）

将特征图划分为不重叠的窗口，仅在窗口内计算自注意力。窗口大小固定为M×M（通常M=7）。

复杂度对比：

• 全局MSA: Ω(MSA) = 4hwC² + 2(hw)²C
• 窗口MSA: Ω(W-MSA) = 4hwC² + 2M²hwC

当hw >> M²时，窗口注意力大幅降低计算量。

(3) 移位窗口注意力（SW-MSA）

W-MSA的问题：窗口之间缺乏信息交互。Swin引入移位窗口机制，在连续的两层中交替使用规则窗口和偏移窗口：

• 第l层：规则窗口划分
• 第l+1层：窗口偏移 (⌊M/2⌋, ⌊M/2⌋)

这样跨窗口的连接在两层之间得以建立，实现了全局建模能力。

高效批处理 ：移位后窗口数量变化，Swin使用cyclic shift 和masked attention实现高效的批处理计算。

3.3 相对位置编码

Swin使用相对位置偏置而非绝对位置编码：

Attention(Q, K, V) = softmax(Q × K^T / √d + B) × V

其中B是相对位置偏置矩阵，可学习参数。相对位置编码对平移更具鲁棒性。

3.4 整体配置（Swin-Tiny为例）

阶段	输出尺寸	层数	通道数	窗口大小
Stage 1	56×56	2	96	7
Stage 2	28×28	2	192	7
Stage 3	14×14	6	384	7
Stage 4	7×7	2	768	7

---

四、模型对比

特性	Transformer	ViT	Swin
输入形式	1D token序列	2D图像patch序列	2D图像patch序列
核心结构	Encoder-Decoder	Encoder Only	Encoder Only
注意力范围	全局	全局	窗口内（+跨窗口）
位置编码	正弦/余弦	可学习绝对位置	可学习相对位置偏置
特征尺度	单尺度	单尺度	多尺度金字塔
主要应用	NLP	图像分类	通用视觉任务
归纳偏置	无	无	局部性（窗口）

---

五、代码实战

5.1 环境准备

pip install torch torchvision timm einops numpy matplotlib

5.2 Transformer核心模块实现

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

class MultiHeadAttention(nn.Module):
    def __init__(self, d_model, n_heads, dropout=0.1):
        super().__init__()
        assert d_model % n_heads == 0
        self.d_model = d_model
        self.n_heads = n_heads
        self.d_k = d_model // n_heads
        
        self.W_q = nn.Linear(d_model, d_model)
        self.W_k = nn.Linear(d_model, d_model)
        self.W_v = nn.Linear(d_model, d_model)
        self.W_o = nn.Linear(d_model, d_model)
        self.dropout = nn.Dropout(dropout)
        
    def forward(self, x, mask=None):
        batch_size, seq_len, _ = x.size()
        
        # 线性投影并分头
        Q = self.W_q(x).view(batch_size, seq_len, self.n_heads, self.d_k).transpose(1, 2)
        K = self.W_k(x).view(batch_size, seq_len, self.n_heads, self.d_k).transpose(1, 2)
        V = self.W_v(x).view(batch_size, seq_len, self.n_heads, self.d_k).transpose(1, 2)
        
        # 计算注意力分数
        scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(self.d_k)
        
        if mask is not None:
            scores = scores.masked_fill(mask == 0, -1e9)
        
        attn = F.softmax(scores, dim=-1)
        attn = self.dropout(attn)
        
        # 加权求和
        out = torch.matmul(attn, V)
        out = out.transpose(1, 2).contiguous().view(batch_size, seq_len, self.d_model)
        out = self.W_o(out)
        return out

class PositionwiseFFN(nn.Module):
    def __init__(self, d_model, d_ff, dropout=0.1):
        super().__init__()
        self.fc1 = nn.Linear(d_model, d_ff)
        self.fc2 = nn.Linear(d_ff, d_model)
        self.dropout = nn.Dropout(dropout)
        self.activation = nn.ReLU()
        
    def forward(self, x):
        return self.fc2(self.dropout(self.activation(self.fc1(x))))

class EncoderLayer(nn.Module):
    def __init__(self, d_model, n_heads, d_ff, dropout=0.1):
        super().__init__()
        self.self_attn = MultiHeadAttention(d_model, n_heads, dropout)
        self.ffn = PositionwiseFFN(d_model, d_ff, dropout)
        self.norm1 = nn.LayerNorm(d_model)
        self.norm2 = nn.LayerNorm(d_model)
        self.dropout = nn.Dropout(dropout)
        
    def forward(self, x, mask=None):
        # 自注意力 + 残差
        attn_out = self.self_attn(x, mask)
        x = x + self.dropout(attn_out)
        x = self.norm1(x)
        
        # FFN + 残差
        ffn_out = self.ffn(x)
        x = x + self.dropout(ffn_out)
        x = self.norm2(x)
        return x

5.3 ViT完整实现

class PatchEmbed(nn.Module):
    def __init__(self, img_size=224, patch_size=16, in_chans=3, embed_dim=768):
        super().__init__()
        self.img_size = img_size
        self.patch_size = patch_size
        self.n_patches = (img_size // patch_size) ** 2
        
        # 使用卷积实现patch embedding
        self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=patch_size)
        
    def forward(self, x):
        x = self.proj(x)  # (B, embed_dim, H/patch, W/patch)
        x = x.flatten(2).transpose(1, 2)  # (B, n_patches, embed_dim)
        return x

class ViT(nn.Module):
    def __init__(
        self,
        img_size=224,
        patch_size=16,
        in_chans=3,
        n_classes=1000,
        embed_dim=768,
        depth=12,
        n_heads=12,
        mlp_ratio=4,
        dropout=0.1,
    ):
        super().__init__()
        
        self.patch_embed = PatchEmbed(img_size, patch_size, in_chans, embed_dim)
        n_patches = self.patch_embed.n_patches
        
        # 可学习的[CLS] token
        self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim))
        # 可学习的位置编码
        self.pos_embed = nn.Parameter(torch.zeros(1, n_patches + 1, embed_dim))
        self.pos_drop = nn.Dropout(dropout)
        
        # Transformer编码器
        self.blocks = nn.ModuleList([
            EncoderLayer(embed_dim, n_heads, int(embed_dim * mlp_ratio), dropout)
            for _ in range(depth)
        ])
        
        self.norm = nn.LayerNorm(embed_dim)
        self.head = nn.Linear(embed_dim, n_classes)
        
        # 初始化
        nn.init.trunc_normal_(self.cls_token, std=0.02)
        nn.init.trunc_normal_(self.pos_embed, std=0.02)
        self.apply(self._init_weights)
        
    def _init_weights(self, m):
        if isinstance(m, nn.Linear):
            nn.init.trunc_normal_(m.weight, std=0.02)
            if m.bias is not None:
                nn.init.zeros_(m.bias)
        elif isinstance(m, nn.LayerNorm):
            nn.init.zeros_(m.bias)
            nn.init.ones_(m.weight)
            
    def forward(self, x):
        B = x.shape[0]
        x = self.patch_embed(x)  # (B, n_patches, embed_dim)
        
        # 添加[CLS] token
        cls_tokens = self.cls_token.expand(B, -1, -1)
        x = torch.cat([cls_tokens, x], dim=1)
        x = x + self.pos_embed
        x = self.pos_drop(x)
        
        # 经过Transformer层
        for block in self.blocks:
            x = block(x)
        
        x = self.norm(x)
        # 取[CLS] token进行分类
        cls_out = x[:, 0]
        return self.head(cls_out)

5.4 Swin Transformer核心组件实现

class WindowAttention(nn.Module):
    def __init__(self, dim, window_size, n_heads, qkv_bias=True, dropout=0.1):
        super().__init__()
        self.dim = dim
        self.window_size = window_size
        self.n_heads = n_heads
        head_dim = dim // n_heads
        self.scale = head_dim ** -0.5
        
        self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
        self.attn_drop = nn.Dropout(dropout)
        self.proj = nn.Linear(dim, dim)
        self.proj_drop = nn.Dropout(dropout)
        
        # 相对位置偏置
        self.relative_position_bias_table = nn.Parameter(
            torch.zeros((2 * window_size[0] - 1) * (2 * window_size[1] - 1), n_heads)
        )
        
        # 计算相对位置索引
        coords_h = torch.arange(window_size[0])
        coords_w = torch.arange(window_size[1])
        coords = torch.stack(torch.meshgrid(coords_h, coords_w, indexing="ij"))
        coords_flatten = coords.flatten(1)  # (2, window_size^2)
        relative_coords = coords_flatten[:, :, None] - coords_flatten[:, None, :]
        relative_coords = relative_coords.permute(1, 2, 0).contiguous()
        relative_coords[:, :, 0] += window_size[0] - 1
        relative_coords[:, :, 1] += window_size[1] - 1
        relative_coords[:, :, 0] *= 2 * window_size[1] - 1
        self.relative_position_index = relative_coords.sum(-1).flatten()
        
        nn.init.trunc_normal_(self.relative_position_bias_table, std=0.02)
        
    def forward(self, x, mask=None):
        B, N, C = x.shape
        qkv = self.qkv(x).reshape(B, N, 3, self.n_heads, C // self.n_heads).permute(2, 0, 3, 1, 4)
        q, k, v = qkv[0], qkv[1], qkv[2]
        
        q = q * self.scale
        attn = (q @ k.transpose(-2, -1))
        
        # 添加相对位置偏置
        relative_position_bias = self.relative_position_bias_table[self.relative_position_index].view(
            self.window_size[0] * self.window_size[1], self.window_size[0] * self.window_size[1], -1
        )
        relative_position_bias = relative_position_bias.permute(2, 0, 1).contiguous()
        attn = attn + relative_position_bias.unsqueeze(0)
        
        if mask is not None:
            nW = mask.shape[0]
            attn = attn.view(B // nW, nW, self.n_heads, N, N) + mask.unsqueeze(1).unsqueeze(0)
            attn = attn.view(-1, self.n_heads, N, N)
            attn = F.softmax(attn, dim=-1)
        else:
            attn = F.softmax(attn, dim=-1)
            
        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

class SwinTransformerBlock(nn.Module):
    def __init__(self, dim, n_heads, window_size=7, shift_size=0, mlp_ratio=4, dropout=0.1):
        super().__init__()
        self.dim = dim
        self.n_heads = n_heads
        self.window_size = window_size
        self.shift_size = shift_size
        self.mlp_ratio = mlp_ratio
        
        self.norm1 = nn.LayerNorm(dim)
        self.attn = WindowAttention(dim, (window_size, window_size), n_heads, dropout=dropout)
        self.norm2 = nn.LayerNorm(dim)
        self.mlp = PositionwiseFFN(dim, int(dim * mlp_ratio), dropout)
        
    def forward(self, x, H, W):
        B, L, C = x.shape
        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 = self._window_partition(shifted_x, self.window_size)
        x_windows = x_windows.view(-1, self.window_size * self.window_size, C)
        
        # 窗口注意力
        attn_windows = self.attn(x_windows, mask=None)  # 实际实现需要处理mask
        attn_windows = attn_windows.view(-1, self.window_size, self.window_size, C)
        
        # 窗口合并
        shifted_x = self._window_reverse(attn_windows, self.window_size, H, W)
        
        # 反向循环移位
        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 = self.norm1(x + x)  # 残差连接（简化）
        
        # MLP
        mlp_out = self.mlp(x)
        x = x + mlp_out
        x = self.norm2(x)
        return x
    
    def _window_partition(self, x, window_size):
        B, H, W, C = x.shape
        x = x.view(B, H // window_size, window_size, W // window_size, window_size, C)
        x = x.permute(0, 1, 3, 2, 4, 5).contiguous()
        return x.view(-1, window_size, window_size, C)
    
    def _window_reverse(self, x, window_size, H, W):
        B = int(x.shape[0] / (H * W / window_size / window_size))
        x = x.view(B, H // window_size, W // window_size, window_size, window_size, -1)
        x = x.permute(0, 1, 3, 2, 4, 5).contiguous()
        return x.view(B, H, W, -1)

5.5 训练与推理代码

import torch.optim as optim
from torch.utils.data import DataLoader
from torchvision import datasets, transforms

def train_one_epoch(model, dataloader, optimizer, criterion, device):
    model.train()
    total_loss = 0
    correct = 0
    total = 0
    
    for batch_idx, (data, target) in enumerate(dataloader):
        data, target = data.to(device), target.to(device)
        optimizer.zero_grad()
        output = model(data)
        loss = criterion(output, target)
        loss.backward()
        optimizer.step()
        
        total_loss += loss.item()
        _, predicted = output.max(1)
        total += target.size(0)
        correct += predicted.eq(target).sum().item()
        
    return total_loss / len(dataloader), 100. * correct / total

def evaluate(model, dataloader, criterion, device):
    model.eval()
    total_loss = 0
    correct = 0
    total = 0
    
    with torch.no_grad():
        for data, target in dataloader:
            data, target = data.to(device), target.to(device)
            output = model(data)
            loss = criterion(output, target)
            total_loss += loss.item()
            _, predicted = output.max(1)
            total += target.size(0)
            correct += predicted.eq(target).sum().item()
            
    return total_loss / len(dataloader), 100. * correct / total

def main():
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    
    # 数据准备
    transform_train = transforms.Compose([
        transforms.RandomResizedCrop(224),
        transforms.RandomHorizontalFlip(),
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
    ])
    
    transform_test = 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])
    ])
    
    # 使用CIFAR-10进行演示（完整训练建议使用ImageNet）
    train_dataset = datasets.CIFAR10(root='./data', train=True, download=True, transform=transform_train)
    test_dataset = datasets.CIFAR10(root='./data', train=False, download=True, transform=transform_test)
    
    train_loader = DataLoader(train_dataset, batch_size=128, shuffle=True, num_workers=4)
    test_loader = DataLoader(test_dataset, batch_size=128, shuffle=False, num_workers=4)
    
    # 初始化模型
    model = ViT(
        img_size=32,  # CIFAR-10尺寸较小
        patch_size=4,
        n_classes=10,
        embed_dim=192,
        depth=6,
        n_heads=6,
    ).to(device)
    
    # 优化器和损失函数
    optimizer = optim.AdamW(model.parameters(), lr=1e-3, weight_decay=0.05)
    scheduler = optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=100)
    criterion = nn.CrossEntropyLoss()
    
    # 训练循环
    n_epochs = 50
    best_acc = 0
    
    for epoch in range(n_epochs):
        train_loss, train_acc = train_one_epoch(model, train_loader, optimizer, criterion, device)
        scheduler.step()
        
        if (epoch + 1) % 5 == 0:
            test_loss, test_acc = evaluate(model, test_loader, criterion, device)
            print(f'Epoch {epoch+1}: Train Loss={train_loss:.4f}, Train Acc={train_acc:.2f}%, Test Acc={test_acc:.2f}%')
            
            if test_acc > best_acc:
                best_acc = test_acc
                torch.save(model.state_dict(), 'best_model.pth')
    
    print(f'Best Test Accuracy: {best_acc:.2f}%')

def inference_example():
    """推理示例：加载模型并对单张图像进行分类"""
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    
    # 加载模型
    model = ViT(img_size=32, patch_size=4, n_classes=10).to(device)
    model.load_state_dict(torch.load('best_model.pth', map_location=device))
    model.eval()
    
    # 图像预处理
    transform = transforms.Compose([
        transforms.Resize(32),
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
    ])
    
    # 假设有一张CIFAR-10图像
    from PIL import Image
    import numpy as np
    
    # 生成随机图像模拟（实际使用中替换为真实图像路径）
    dummy_img = Image.fromarray(np.random.randint(0, 255, (32, 32, 3), dtype=np.uint8))
    img_tensor = transform(dummy_img).unsqueeze(0).to(device)
    
    with torch.no_grad():
        output = model(img_tensor)
        prob = F.softmax(output, dim=1)
        pred_class = output.argmax(dim=1).item()
        
    cifar10_classes = ['airplane', 'automobile', 'bird', 'cat', 'deer', 
                       'dog', 'frog', 'horse', 'ship', 'truck']
    print(f'Predicted class: {cifar10_classes[pred_class]}')
    print(f'Confidence: {prob[0][pred_class].item():.4f}')

if __name__ == '__main__':
    # main()  # 训练
    inference_example()  # 推理

六、实践建议与常见问题

6.1 训练技巧

1. 学习率调度：使用warmup（前几个epoch线性增加学习率）+ cosine decay是ViT和Swin的标准配置
2. 数据增强：ViT和Swin对数据增强敏感，推荐使用RandAugment或AutoAugment
3. 权重初始化：使用truncated normal初始化，标准差0.02
4. Batch Size：尽量大（推荐1024+），配合梯度累积

6.2 常见问题及解决

问题	可能原因	解决方案
训练不收敛	学习率过大/过小	调整LR，使用warmup
过拟合	模型大，数据少	增加数据增强，提高dropout
内存不足	序列太长	减小batch size，使用gradient checkpointing
分类精度低	预训练不足	使用ImageNet预训练权重微调

6.3 模型选择指南

1. 小数据集（<100k）：建议使用CNN（ResNet）或Swin（利用其局部性）
2. 中等数据集（100k-1M）：ViT或Swin均可，Swin通常更好
3. 大数据集（>1M）：ViT表现优异，扩展性好
4. 检测/分割任务：Swin的层级特征天然适配
5. 计算资源有限：Swin-Tiny或ViT-Small