Vision Transformer (ViT) 从零实现：PyTorch 代码逐行解析 3 大核心模块

Vision Transformer (ViT) 从零实现：PyTorch 代码逐行解析 3 大核心模块

当卷积神经网络（CNN）长期主导计算机视觉领域时，Google Research在2020年提出的Vision Transformer（ViT）彻底改变了这一格局。本文将带您从零开始构建一个完整的ViT-B/16模型，通过PyTorch代码逐行解析其三大核心模块：Patch Embedding、Transformer Encoder和MLP Head。不同于理论概述，我们聚焦于工程实现细节，让您不仅能理解原理，更能亲手实现这一革命性架构。

1. 环境准备与整体架构

在开始编码之前，我们需要配置开发环境并理解ViT的整体架构设计。以下是推荐的开发环境配置：

import torch import torch.nn as nn import torch.nn.functional as F from torch import einsum from einops import rearrange, repeat print(f"PyTorch版本: {torch.__version__}") # 输出示例: PyTorch版本: 2.0.1

ViT-B/16的主要参数配置如下表所示：

提示：实际实现时建议将这些参数封装为配置类，方便后续调整和实验管理。

ViT的工作流程可以概括为以下三个阶段：

1. 图像分块嵌入：将2D图像转换为1D序列
2. Transformer编码：通过自注意力机制提取全局特征
3. 分类头处理：基于[CLS]token输出分类结果

2. Patch Embedding实现详解

Patch Embedding模块负责将2D图像转换为Transformer可处理的1D序列，这是ViT区别于CNN的关键设计。其核心思想是将图像分割为固定大小的块，然后线性投影到嵌入空间。

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, img_size) self.patch_size = (patch_size, patch_size) self.num_patches = (img_size // patch_size) ** 2 # 使用卷积层实现分块和投影 self.proj = nn.Conv2d( in_chans, embed_dim, kernel_size=patch_size, stride=patch_size ) # 可学习的分类token和位置编码 self.cls_token = nn.Parameter(torch.randn(1, 1, embed_dim)) self.pos_embed = nn.Parameter( torch.randn(1, self.num_patches + 1, embed_dim) ) def forward(self, x): B, C, H, W = x.shape assert H == self.img_size[0] and W == self.img_size[1], \ f"输入图像尺寸({H}*{W})与模型设置({self.img_size[0]}*{self.img_size[1]})不符" # 分块投影 [B, C, H, W] -> [B, embed_dim, num_patches_w, num_patches_h] x = self.proj(x) # 展平并转置 [B, embed_dim, num_patches_w, num_patches_h] -> [B, num_patches, embed_dim] x = x.flatten(2).transpose(1, 2) # 添加分类token [B, 1, embed_dim] cls_tokens = self.cls_token.expand(B, -1, -1) x = torch.cat((cls_tokens, x), dim=1) # 添加位置编码 x = x + self.pos_embed return x

这段代码实现了几个关键功能：

• 使用卷积核和步长等于patch_size的Conv2d层同时完成分块和线性投影
• 通过可学习的cls_token实现类似BERT的[CLS]标记
• 添加可学习的位置编码保留空间信息

注意：原始论文发现1D位置编码效果与更复杂的2D编码相当，因此这里采用简单的可学习1D编码。

3. Transformer Encoder深度解析

Transformer Encoder是ViT的核心特征提取器，由多个相同的编码器层堆叠而成。每个编码器层包含多头自注意力机制（MSA）和前馈网络（MLP）。

3.1 多头自注意力实现

class MultiHeadAttention(nn.Module): def __init__(self, embed_dim=768, num_heads=12): super().__init__() self.embed_dim = embed_dim self.num_heads = num_heads self.head_dim = embed_dim // num_heads self.qkv = nn.Linear(embed_dim, embed_dim * 3) self.proj = nn.Linear(embed_dim, embed_dim) def forward(self, x): B, N, C = x.shape # 生成QKV [B, N, 3*embed_dim] qkv = self.qkv(x) # 分离QKV并分头 [B, num_heads, N, head_dim] q, k, v = rearrange( qkv, 'b n (qkv h d) -> qkv b h n d', qkv=3, h=self.num_heads, d=self.head_dim ) # 缩放点积注意力 attn = (q @ k.transpose(-2, -1)) * (self.head_dim ** -0.5) attn = attn.softmax(dim=-1) # 加权求和 [B, h, N, d] out = attn @ v # 合并多头 [B, N, embed_dim] out = rearrange(out, 'b h n d -> b n (h d)') out = self.proj(out) return out

3.2 编码器层完整实现

class TransformerEncoderLayer(nn.Module): def __init__(self, embed_dim=768, num_heads=12, mlp_ratio=4.0): super().__init__() self.norm1 = nn.LayerNorm(embed_dim) self.attn = MultiHeadAttention(embed_dim, num_heads) self.norm2 = nn.LayerNorm(embed_dim) mlp_hidden_dim = int(embed_dim * mlp_ratio) self.mlp = nn.Sequential( nn.Linear(embed_dim, mlp_hidden_dim), nn.GELU(), nn.Linear(mlp_hidden_dim, embed_dim), ) def forward(self, x): # 残差连接1 x = x + self.attn(self.norm1(x)) # 残差连接2 x = x + self.mlp(self.norm2(x)) return x

编码器层的关键设计特点：

• 前置LayerNorm：与原始Transformer不同，ViT采用前置归一化
• GELU激活：相比ReLU，GELU在Transformer中表现更好
• 残差连接：确保梯度有效传播，支持深层网络训练

4. MLP Head与完整模型集成

MLP Head负责将Transformer输出的[CLS]token特征映射到分类空间。虽然结构简单，但有一些实现细节需要注意。

class VisionTransformer(nn.Module): def __init__(self, img_size=224, patch_size=16, in_chans=3, num_classes=1000, embed_dim=768, depth=12, num_heads=12, mlp_ratio=4.0): super().__init__() self.patch_embed = PatchEmbedding( img_size, patch_size, in_chans, embed_dim ) # 构建Transformer编码器 self.blocks = nn.ModuleList([ TransformerEncoderLayer(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): # 分块嵌入 x = self.patch_embed(x) # 通过所有编码器层 for blk in self.blocks: x = blk(x) # 最终归一化 x = self.norm(x) # 取[CLS]token进行分类 cls_token = x[:, 0] return self.head(cls_token)

完整模型的关键集成点：

1. 模块组合：将Patch Embedding、Transformer Encoder和MLP Head串联
2. 分类策略：仅使用[CLS]token对应的特征进行分类
3. 归一化位置：在最后一层后添加额外的LayerNorm

5. 训练技巧与实战建议

实现模型只是第一步，要让ViT真正work还需要注意以下训练细节：

学习率设置

ViT对学习率非常敏感，推荐使用带warmup的余弦退火调度：

def get_optimizer(model, lr=3e-5, warmup_steps=10000): optimizer = torch.optim.AdamW(model.parameters(), lr=lr) scheduler = torch.optim.lr_scheduler.LambdaLR( optimizer, lr_lambda=lambda step: min( (step + 1) / (warmup_steps + 1), 1.0 / (step ** 0.5) ) ) return optimizer, scheduler

数据增强策略

由于ViT缺乏CNN的归纳偏置，需要更强的数据增强：

from torchvision import transforms train_transform = transforms.Compose([ transforms.RandomResizedCrop(224, scale=(0.08, 1.0)), transforms.RandomHorizontalFlip(), transforms.ColorJitter(0.4, 0.4, 0.4), transforms.ToTensor(), transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]) ])

混合精度训练

使用AMP加速训练并减少显存占用：

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

6. 模型变体与扩展应用

基于ViT的基础架构，研究者们提出了多种改进版本。下表对比了几种常见变体：

这些变体主要在以下方面进行了优化：

• 计算效率：如Swin的局部注意力窗口
• 数据效率：如DeiT的蒸馏策略
• 多尺度处理：如CrossViT的双分支设计

7. 部署优化与性能调优

在实际部署ViT模型时，可以考虑以下优化技术：

模型量化

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

ONNX导出

dummy_input = torch.randn(1, 3, 224, 224) torch.onnx.export( model, dummy_input, "vit.onnx", input_names=["input"], output_names=["output"], dynamic_axes={"input": {0: "batch"}, "output": {0: "batch"}} )

TensorRT加速

# 使用torch2trt等工具转换模型 from torch2trt import torch2trt model_trt = torch2trt( model, [dummy_input], fp16_mode=True, max_workspace_size=1<<25 )