Vision Transformer实战：从零开始用PyTorch搭建ViT模型（附完整代码）

Vision Transformer实战：从零搭建ViT模型与工业级优化技巧

1. 环境准备与数据预处理

在开始构建ViT模型之前，我们需要搭建合适的开发环境并准备图像数据。与传统的CNN不同，ViT对输入数据的处理有独特要求，这直接影响到模型的最终性能。

推荐开发环境配置：

conda create -n vit_env python=3.8 conda activate vit_env pip install torch==1.12.0+cu113 torchvision==0.13.0+cu113 -f https://download.pytorch.org/whl/torch_stable.html pip install timm==0.6.7 matplotlib pandas

对于图像数据处理，ViT需要将图像分割为固定大小的patch。以下是关键的预处理步骤：

1. 图像尺寸标准化：将所有输入图像调整为统一尺寸（通常为224×224或384×384）
2. Patch分割：将图像划分为N×N的patch（常用16×16或32×32）
3. 归一化处理：应用ImageNet标准的均值和标准差进行归一化

from torchvision import transforms # ViT标准数据增强流程 train_transform = 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]) ]) 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]) ])

注意：对于高分辨率任务（如医疗影像），可考虑增大patch尺寸以减少计算量，但会损失细粒度信息

2. ViT模型架构深度解析

2.1 Patch Embedding层实现

Patch Embedding是ViT区别于CNN的核心组件，它将图像转换为Transformer可处理的序列形式。以下是关键实现细节：

import torch import torch.nn as nn class PatchEmbedding(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 self.proj = nn.Conv2d( in_chans, embed_dim, kernel_size=patch_size, stride=patch_size ) # 使用卷积操作实现patch分割和投影 def forward(self, x): x = self.proj(x) # (B, E, H/P, W/P) x = x.flatten(2) # (B, E, N) x = x.transpose(1, 2) # (B, N, E) return x

参数选择对比表：

2.2 Transformer Encoder设计

ViT的核心是由多个Transformer Encoder层堆叠而成。每个Encoder包含以下组件：

1. 多头注意力机制：计算patch间的全局关系
2. MLP块：特征非线性变换
3. LayerNorm：稳定训练过程
4. 残差连接：缓解梯度消失

class TransformerBlock(nn.Module): def __init__(self, embed_dim, num_heads, mlp_ratio=4.0, dropout=0.1): super().__init__() self.norm1 = nn.LayerNorm(embed_dim) self.attn = nn.MultiheadAttention(embed_dim, num_heads, dropout=dropout) self.norm2 = nn.LayerNorm(embed_dim) self.mlp = nn.Sequential( nn.Linear(embed_dim, int(embed_dim * mlp_ratio)), nn.GELU(), nn.Dropout(dropout), nn.Linear(int(embed_dim * mlp_ratio), embed_dim), nn.Dropout(dropout) ) def forward(self, x): # 注意力部分 res = x x = self.norm1(x) x, _ = self.attn(x, x, x) x = res + x # MLP部分 res = x x = self.norm2(x) x = self.mlp(x) x = res + x return x

3. 完整ViT模型实现

结合上述组件，我们可以构建完整的ViT模型：

class VisionTransformer(nn.Module): def __init__(self, img_size=224, patch_size=16, in_chans=3, embed_dim=768, depth=12, num_heads=12, mlp_ratio=4., num_classes=1000): super().__init__() # Patch嵌入 self.patch_embed = PatchEmbedding(img_size, patch_size, in_chans, embed_dim) # 分类token和位置编码 self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim)) self.pos_embed = nn.Parameter( torch.zeros(1, self.patch_embed.n_patches + 1, embed_dim) ) # Transformer编码器 self.blocks = nn.ModuleList([ TransformerBlock(embed_dim, num_heads, mlp_ratio) for _ in range(depth) ]) # 分类头 self.norm = nn.LayerNorm(embed_dim) self.head = nn.Linear(embed_dim, num_classes) def forward(self, x): B = x.shape[0] # 生成patch嵌入 x = self.patch_embed(x) # (B, N, E) # 添加分类token cls_tokens = self.cls_token.expand(B, -1, -1) x = torch.cat((cls_tokens, x), dim=1) # 添加位置编码 x = x + self.pos_embed # 通过Transformer编码器 for block in self.blocks: x = block(x) # 分类 x = self.norm(x) cls_token_final = x[:, 0] x = self.head(cls_token_final) return x

4. 训练技巧与性能优化

4.1 学习率调度策略

ViT训练对学习率非常敏感，推荐采用warmup+cosine衰减策略：

from torch.optim import AdamW from torch.optim.lr_scheduler import CosineAnnealingLR optimizer = AdamW(model.parameters(), lr=1e-4, weight_decay=0.05) scheduler = CosineAnnealingLR(optimizer, T_max=epochs, eta_min=1e-6) # Warmup实现 def warmup_lr_scheduler(optimizer, warmup_iters, warmup_factor): def f(x): if x >= warmup_iters: return 1 alpha = float(x) / warmup_iters return warmup_factor * (1 - alpha) + alpha return torch.optim.lr_scheduler.LambdaLR(optimizer, f)

4.2 混合精度训练

使用AMP(自动混合精度)可显著减少显存占用并加速训练：

from torch.cuda.amp import GradScaler, autocast scaler = GradScaler() for inputs, targets in train_loader: optimizer.zero_grad() with autocast(): outputs = model(inputs) loss = criterion(outputs, targets) scaler.scale(loss).backward() scaler.step(optimizer) scaler.update()

4.3 关键超参数设置

基于实验经验的参数推荐：

5. 模型微调与部署实践

5.1 迁移学习技巧

当在特定领域数据上微调ViT时：

1. 分层学习率：不同层使用不同学习率

param_groups = [ {'params': model.patch_embed.parameters(), 'lr': base_lr*0.1}, {'params': model.pos_embed, 'lr': base_lr*0.5}, {'params': model.cls_token, 'lr': base_lr}, {'params': model.blocks.parameters(), 'lr': base_lr}, {'params': model.head.parameters(), 'lr': base_lr*2} ]

2. 渐进式解冻：从顶层开始逐步解冻底层参数

5.2 模型量化部署

使用TorchScript量化减小模型体积：

quantized_model = torch.quantization.quantize_dynamic( model, {torch.nn.Linear}, dtype=torch.qint8 ) traced_script = torch.jit.trace(quantized_model, example_input) traced_script.save("vit_quantized.pt")

部署性能对比：

在实际项目中，ViT模型经过适当优化后，在ImageNet-1k上可以达到约80%的top-1准确率，同时保持合理的计算效率。相比传统CNN，ViT在数据充足时展现出更强的表征能力，特别适合需要全局上下文理解的任务。