Vision Transformer (ViT) 代码实现详解：从 Patch Embedding 到 MLP Head 的 5 个关键模块

当计算机视觉领域还在被卷积神经网络（CNN）统治时，Google Research 在 2020 年提出的 Vision Transformer (ViT) 彻底改变了游戏规则。这篇技术指南将带您深入 ViT 的 PyTorch 实现细节，通过 5 个核心模块的代码级解析，让您不仅能理解 ViT 的工作原理，还能亲手实现这个革命性的架构。

1. Patch Embedding：图像到序列的魔法转换

传统 Transformer 处理的是序列数据，而图像本质上是二维的。ViT 的第一个关键创新就是将图像转换为适合 Transformer 处理的序列形式。让我们看看如何用 PyTorch 实现这一转换：

import torch import torch.nn as nn class PatchEmbedding(nn.Module): def __init__(self, img_size=224, patch_size=16, in_channels=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分割和线性投影 self.proj = nn.Conv2d( in_channels=in_channels, out_channels=embed_dim, kernel_size=patch_size, stride=patch_size ) def forward(self, x): # x形状: [batch_size, channels, height, width] x = self.proj(x) # [batch_size, embed_dim, n_patches^0.5, n_patches^0.5] x = x.flatten(2) # [batch_size, embed_dim, n_patches] x = x.transpose(1, 2) # [batch_size, n_patches, embed_dim] return x

这段代码有几个精妙之处：

1. 卷积核与步长的设计：kernel_size=patch_size和stride=patch_size确保每个卷积操作正好处理一个 patch，没有重叠
2. 维度变换的艺术：通过flatten和transpose操作，我们将 2D 特征图转换为适合 Transformer 处理的序列形式
3. 计算效率：相比先分割图像再线性投影，这种实现方式更加高效

提示：对于 224x224 的图像和 16x16 的 patch，会产生 196 个 patch（(224/16)²）。每个 patch 被投影到 768 维空间（ViT-Base 的默认设置）。

2. Position Embedding：为序列注入空间信息

Transformer 本身不具备处理序列顺序的能力，而图像中的空间位置信息至关重要。ViT 使用可学习的位置编码来解决这个问题：

class PositionEmbedding(nn.Module): def __init__(self, n_patches=196, embed_dim=768): super().__init__() self.pos_embed = nn.Parameter(torch.zeros(1, n_patches + 1, embed_dim)) nn.init.trunc_normal_(self.pos_embed, std=0.02) def forward(self, x): # x形状: [batch_size, n_patches, embed_dim] # 添加位置编码（不包括class token的位置） return x + self.pos_embed[:, 1:, :]

值得注意的是：

1. 可学习参数：与原始 Transformer 的固定位置编码不同，ViT 使用可学习的位置编码
2. class token 处理：位置编码的第一个位置保留给 class token
3. 初始化技巧：使用截断正态分布初始化，标准差为 0.02

3. Transformer Encoder Block：自注意力的核心

ViT 的编码器由多个相同的 Transformer Block 堆叠而成，每个 Block 包含多头自注意力机制和前馈网络：

class TransformerBlock(nn.Module): def __init__(self, embed_dim=768, num_heads=12, 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): # LayerNorm前置（Pre-LN）结构 norm_x = self.norm1(x) attn_output, _ = self.attn(norm_x, norm_x, norm_x) x = x + attn_output norm_x = self.norm2(x) mlp_output = self.mlp(norm_x) x = x + mlp_output return x

关键实现细节：

1. Pre-LN 结构：LayerNorm 在自注意力和 MLP 之前应用，相比原始 Transformer 的 Post-LN 更稳定
2. 残差连接：每个子层都有残差连接，缓解梯度消失问题
3. GELU 激活：使用高斯误差线性单元，比 ReLU 表现更好

4. MLP Head：从特征到分类结果

ViT 使用一个简单的 MLP 作为分类头，处理来自 class token 的特征：

class MLPHead(nn.Module): def __init__(self, embed_dim=768, num_classes=1000): super().__init__() self.head = nn.Sequential( nn.LayerNorm(embed_dim), nn.Linear(embed_dim, num_classes) ) def forward(self, x): # x形状: [batch_size, embed_dim] (class token) return self.head(x)

这个看似简单的结构有几个设计考量：

1. 仅使用 class token：与 CNN 通常使用全局平均池化不同，ViT 依赖 class token 的特征
2. 最后的 LayerNorm：确保特征稳定，有助于模型收敛
3. 线性分类器：微调时可以使用更复杂的结构，但预训练时保持简单

5. 完整 ViT 模型集成

现在我们将所有模块组合成完整的 ViT 模型：

class VisionTransformer(nn.Module): def __init__(self, img_size=224, patch_size=16, in_channels=3, embed_dim=768, depth=12, num_heads=12, mlp_ratio=4.0, num_classes=1000, dropout=0.1): super().__init__() self.patch_embed = PatchEmbedding(img_size, patch_size, in_channels, embed_dim) n_patches = self.patch_embed.n_patches # Class token self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim)) # Position embedding self.pos_embed = PositionEmbedding(n_patches, embed_dim) # Transformer encoder self.blocks = nn.ModuleList([ TransformerBlock(embed_dim, num_heads, mlp_ratio, dropout) for _ in range(depth) ]) # Classification head self.head = MLPHead(embed_dim, num_classes) # Initialize weights nn.init.trunc_normal_(self.cls_token, 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) def forward(self, x): # Patch embedding x = self.patch_embed(x) # [batch_size, n_patches, embed_dim] # Add class token cls_tokens = self.cls_token.expand(x.shape[0], -1, -1) x = torch.cat((cls_tokens, x), dim=1) # [batch_size, n_patches+1, embed_dim] # Add position embedding x = x + self.pos_embed(x) # Transformer blocks for block in self.blocks: x = block(x) # Classification (use only class token) cls_token_final = x[:, 0] return self.head(cls_token_final)

完整模型的几个关键点：

1. 模块化设计：清晰分离各个组件，便于理解和修改
2. 参数初始化：使用截断正态分布初始化，对 Transformer 模型很重要
3. 前向传播流程：严格遵循 ViT 论文中的数据处理流程

实战：训练 ViT 的实用技巧

理解了 ViT 的实现后，让我们看看如何有效地训练它：

1. 学习率调度：使用带 warmup 的余弦退火调度

   optimizer = torch.optim.AdamW(model.parameters(), lr=1e-4) scheduler = torch.optim.lr_scheduler.CosineAnnealingLR( optimizer, T_max=epochs, eta_min=1e-6 ) warmup_scheduler = torch.optim.lr_scheduler.LinearLR( optimizer, start_factor=0.01, total_iters=warmup_epochs )

2. 数据增强：对 ViT 特别重要

   from torchvision import transforms train_transform = transforms.Compose([ transforms.RandomResizedCrop(224), transforms.RandomHorizontalFlip(), transforms.ColorJitter(brightness=0.2, contrast=0.2, saturation=0.2), transforms.ToTensor(), transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) ])

3. 混合精度训练：加速训练并减少内存使用

   scaler = torch.cuda.amp.GradScaler() with torch.cuda.amp.autocast(): outputs = model(inputs) loss = criterion(outputs, labels) scaler.scale(loss).backward() scaler.step(optimizer) scaler.update()

ViT 变体与性能对比

不同规模的 ViT 模型在性能和计算成本上有显著差异：

注意：这些结果是在 ImageNet-21k 上预训练并在 ImageNet-1k 上微调得到的。实际性能会因训练设置和数据增强策略而有所不同。

常见问题与解决方案

在实现和训练 ViT 时，可能会遇到以下问题：

1. 训练不稳定

   • 解决方案：使用梯度裁剪（nn.utils.clip_grad_norm_）和学习率 warmup

2. 过拟合

   • 解决方案：增加 dropout 率、使用更强的数据增强、添加权重衰减

3. 内存不足

   • 解决方案：减小 batch size、使用梯度累积、尝试混合精度训练

4. 收敛速度慢

   • 解决方案：检查初始化、使用适当的学习率调度、确保数据预处理正确

进阶技巧与优化

对于希望进一步提升 ViT 性能的开发者，可以考虑以下优化：

1. 知识蒸馏：使用更大的 ViT 或 CNN 作为教师模型

   teacher_model = torch.hub.load('facebookresearch/deit:main', 'deit_base_patch16_224', pretrained=True)

2. 渐进式训练：从小分辨率开始，逐步增加图像尺寸

3. 模型剪枝：移除不重要的注意力头或 MLP 神经元

4. 量化感知训练：为部署准备低精度模型

model = torch.quantization.quantize_dynamic( model, {nn.Linear}, dtype=torch.qint8 )

从理论到实践：ViT 的完整训练流程

为了帮助您快速上手，以下是训练 ViT 的完整 PyTorch Lightning 示例：

import pytorch_lightning as pl from torchmetrics import Accuracy class ViTLightningModule(pl.LightningModule): def __init__(self, num_classes=1000, lr=1e-4, warmup_epochs=5): super().__init__() self.model = VisionTransformer(num_classes=num_classes) self.criterion = nn.CrossEntropyLoss() self.train_acc = Accuracy(task='multiclass', num_classes=num_classes) self.val_acc = Accuracy(task='multiclass', num_classes=num_classes) self.lr = lr self.warmup_epochs = warmup_epochs def forward(self, x): return self.model(x) def training_step(self, batch, batch_idx): x, y = batch logits = self(x) loss = self.criterion(logits, y) self.train_acc(logits, y) self.log('train_loss', loss, prog_bar=True) self.log('train_acc', self.train_acc, prog_bar=True) return loss def validation_step(self, batch, batch_idx): x, y = batch logits = self(x) loss = self.criterion(logits, y) self.val_acc(logits, y) self.log('val_loss', loss, prog_bar=True) self.log('val_acc', self.val_acc, prog_bar=True) def configure_optimizers(self): optimizer = torch.optim.AdamW(self.parameters(), lr=self.lr) scheduler = { 'scheduler': torch.optim.lr_scheduler.CosineAnnealingLR( optimizer, T_max=self.trainer.max_epochs, eta_min=1e-6 ), 'interval': 'epoch' } if self.warmup_epochs > 0: warmup = { 'scheduler': torch.optim.lr_scheduler.LinearLR( optimizer, start_factor=0.01, total_iters=self.warmup_epochs ), 'interval': 'epoch' } return [optimizer], [warmup, scheduler] return [optimizer], [scheduler]

这个实现包含了：

1. 训练和验证逻辑：清晰的步骤分离
2. 指标跟踪：准确率和损失记录
3. 学习率调度：带 warmup 的余弦退火
4. 日志记录：实时监控训练进度

ViT 在实际项目中的应用建议

根据实际项目经验，以下是在不同场景下使用 ViT 的建议：

1. 数据量充足时：直接使用 ViT 进行端到端训练
2. 数据有限时：
   • 使用预训练 ViT 作为特征提取器
   • 在顶层添加任务特定层进行微调
3. 计算资源受限时：
   • 考虑较小的 ViT 变体（如 ViT-Tiny 或 ViT-Small）
   • 使用知识蒸馏训练更小的学生模型
4. 实时性要求高时：
   • 探索 ViT 的轻量级变体（如 DeiT 或 MobileViT）
   • 应用模型量化和剪枝

性能优化与调试技巧

当 ViT 表现不如预期时，可以尝试以下调试方法：

1. 可视化注意力图：检查模型关注了图像的哪些部分

   attention_maps = model.blocks[-1].attn.get_attention_map()

2. 检查梯度流动：使用torchviz可视化计算图

   from torchviz import make_dot make_dot(loss, params=dict(model.named_parameters()))

3. 监控激活统计：确保没有异常大的激活值

   for name, param in model.named_parameters(): if param.requires_grad: self.logger.experiment.add_histogram(name, param, self.global_step)

4. 对比基线：与 CNN 基准模型（如 ResNet）比较性能

未来方向与扩展阅读

ViT 开启了计算机视觉的新时代，以下是有前景的研究方向：

1. 高效注意力机制：如 Linformer 或 Performer 中的线性注意力
2. 层次化 ViT：结合 CNN 的层次化特征提取
3. 自监督学习：如 MAE（Masked Autoencoder）方法
4. 多模态应用：如 CLIP 中的图像-文本联合训练

推荐阅读的进阶论文：

• DeiT: Data-efficient Image Transformers
• Swin Transformer: Hierarchical Vision Transformer
• MAE: Masked Autoencoders Are Scalable Vision Learners
• MobileViT: Light-weight, General-purpose, and Mobile-friendly Vision Transformer

结语：ViT 实现的工程实践

实现 ViT 不仅仅是复制论文中的公式，更需要理解其设计哲学和工程细节。在实际项目中，我发现以下几点特别重要：

1. 初始化策略：正确的初始化对 Transformer 的稳定训练至关重要
2. 学习率调度：warmup 阶段能显著提高训练稳定性
3. 正则化技术：适当的 dropout 和权重衰减防止过拟合
4. 硬件利用：充分利用混合精度训练和分布式训练加速实验迭代

ViT 的成功证明了 Transformer 在计算机视觉中的巨大潜力。通过本指南中的代码实现和技巧，希望您能快速掌握这一强大工具，并在自己的项目中应用它。记住，理解每个模块的设计动机比简单复制代码更重要，这将帮助您灵活应对不同的应用场景和挑战。