从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 shiftmasked 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 环境准备

bash 复制代码
pip install torch torchvision timm einops numpy matplotlib

5.2 Transformer核心模块实现

python 复制代码
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完整实现

python 复制代码
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核心组件实现

python 复制代码
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 训练与推理代码

python 复制代码
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
相关推荐
新知图书17 小时前
详细实现与核心配置(新闻早报智能体开发)
人工智能·agent·ai agent·智能体·扣子
Lorin 洛林17 小时前
OpenClaw:跳出聊天框,拥有自主行动力的 AI 智能体
人工智能
rain_sxr17 小时前
大规模数据可视化渲染:D3 与 ECharts 的管线优化实战
人工智能
大海变好AI17 小时前
AIGC分层推理落地能否串联多工具完成自动化闭环
人工智能·python·自动化·aigc
Cosolar17 小时前
AI Agent 架构原理详解:从一次提问到任务完成的完整闭环
人工智能·后端·架构
小白说大模型17 小时前
从向量嵌入到复杂 Agent:LLM、LangChain、LangGraph 完整科普
java·开发语言·人工智能·gpt·深度学习·langchain
协享科技17 小时前
2026 年 AI 表格工具怎么选:ChatGPT、原生 Copilot/Gemini 与公式助手
人工智能·chatgpt·excel·copilot
geo搜搜果数据17 小时前
实测AI搜索GEO监测工具:对比DeepSeek与豆包品牌排名差异
人工智能·langchain·embedding·搜搜果
王莎莎18 小时前
MCP 解决的是工具接入,科研 Agent 还缺的是科学证据接口标准化
前端·人工智能