从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为例):
- 输入图像: 224×224×3
- 分块: 16×16的patch,得到14×14=196个patch
- 线性投影: 每个patch展平为16×16×3=768维,通过线性层映射到D维
- 添加位置编码: 可学习的位置编码
- 添加CLS token: 用于分类的额外可学习token
- 经过L层Transformer Encoder(L=12 for ViT-B)
- 分类头: 取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虽然在分类任务上表现优异,但作为通用视觉骨干存在两个问题:
- 计算复杂度:与图像尺寸平方相关,难以处理高分辨率
- 缺乏多尺度特征:不适合检测、分割等需要多尺度信息的任务
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 环境准备
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 训练技巧
- 学习率调度:使用warmup(前几个epoch线性增加学习率)+ cosine decay是ViT和Swin的标准配置
- 数据增强:ViT和Swin对数据增强敏感,推荐使用RandAugment或AutoAugment
- 权重初始化:使用truncated normal初始化,标准差0.02
- Batch Size:尽量大(推荐1024+),配合梯度累积
6.2 常见问题及解决
| 问题 | 可能原因 | 解决方案 |
|---|---|---|
| 训练不收敛 | 学习率过大/过小 | 调整LR,使用warmup |
| 过拟合 | 模型大,数据少 | 增加数据增强,提高dropout |
| 内存不足 | 序列太长 | 减小batch size,使用gradient checkpointing |
| 分类精度低 | 预训练不足 | 使用ImageNet预训练权重微调 |
6.3 模型选择指南
- 小数据集(<100k):建议使用CNN(ResNet)或Swin(利用其局部性)
- 中等数据集(100k-1M):ViT或Swin均可,Swin通常更好
- 大数据集(>1M):ViT表现优异,扩展性好
- 检测/分割任务:Swin的层级特征天然适配
- 计算资源有限:Swin-Tiny或ViT-Small