Vision Transformer (ViT) 从零实现PyTorch 代码逐行解析 3 大核心模块当卷积神经网络CNN长期主导计算机视觉领域时Google Research在2020年提出的Vision TransformerViT彻底改变了这一格局。本文将带您从零开始构建一个完整的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(fPyTorch版本: {torch.__version__}) # 输出示例: PyTorch版本: 2.0.1ViT-B/16的主要参数配置如下表所示参数名称值说明image_size224输入图像分辨率patch_size16图像分块大小num_channels3输入图像通道数embed_dim768嵌入维度Dnum_heads12注意力头数num_layers12Transformer编码器层数mlp_ratio4.0MLP扩展系数num_classes1000分类类别数ImageNet-1k提示实际实现时建议将这些参数封装为配置类方便后续调整和实验管理。ViT的工作流程可以概括为以下三个阶段图像分块嵌入将2D图像转换为1D序列Transformer编码通过自注意力机制提取全局特征分类头处理基于[CLS]token输出分类结果2. Patch Embedding实现详解Patch Embedding模块负责将2D图像转换为Transformer可处理的1D序列这是ViT区别于CNN的关键设计。其核心思想是将图像分割为固定大小的块然后线性投影到嵌入空间。class PatchEmbedding(nn.Module): def __init__(self, img_size224, patch_size16, in_chans3, embed_dim768): 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_sizepatch_size, stridepatch_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), dim1) # 添加位置编码 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_dim768, num_heads12): 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, qkv3, hself.num_heads, dself.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 out3.2 编码器层完整实现class TransformerEncoderLayer(nn.Module): def __init__(self, embed_dim768, num_heads12, mlp_ratio4.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激活相比ReLUGELU在Transformer中表现更好残差连接确保梯度有效传播支持深层网络训练4. MLP Head与完整模型集成MLP Head负责将Transformer输出的[CLS]token特征映射到分类空间。虽然结构简单但有一些实现细节需要注意。class VisionTransformer(nn.Module): def __init__(self, img_size224, patch_size16, in_chans3, num_classes1000, embed_dim768, depth12, num_heads12, mlp_ratio4.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)完整模型的关键集成点模块组合将Patch Embedding、Transformer Encoder和MLP Head串联分类策略仅使用[CLS]token对应的特征进行分类归一化位置在最后一层后添加额外的LayerNorm5. 训练技巧与实战建议实现模型只是第一步要让ViT真正work还需要注意以下训练细节学习率设置ViT对学习率非常敏感推荐使用带warmup的余弦退火调度def get_optimizer(model, lr3e-5, warmup_steps10000): optimizer torch.optim.AdamW(model.parameters(), lrlr) scheduler torch.optim.lr_scheduler.LambdaLR( optimizer, lr_lambdalambda 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的基础架构研究者们提出了多种改进版本。下表对比了几种常见变体模型变体核心改进参数量ImageNet Top-1ViT-Base原始基准模型86M77.9%DeiT知识蒸馏数据高效训练86M83.1%Swin Transformer分层特征滑动窗口注意力88M83.5%T2T-ViT渐进式token聚合39M82.5%CrossViT多尺度patch融合73M82.8%这些变体主要在以下方面进行了优化计算效率如Swin的局部注意力窗口数据效率如DeiT的蒸馏策略多尺度处理如CrossViT的双分支设计7. 部署优化与性能调优在实际部署ViT模型时可以考虑以下优化技术模型量化quantized_model torch.quantization.quantize_dynamic( model, {torch.nn.Linear}, dtypetorch.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_modeTrue, max_workspace_size125 )