 代码复现：PyTorch 实现 224x224 图像 16x16 Patch 嵌入)
Vision Transformer (ViT) 代码实战从零构建224x224图像的16x16 Patch嵌入与Transformer编码器当计算机视觉领域长期被卷积神经网络CNN统治时Google Research在2020年提出的Vision TransformerViT如同一记惊雷证明了纯Transformer架构在图像识别任务中的惊人潜力。本文将带您深入ViT的核心实现细节通过PyTorch从零构建完整的ViT-B/16模型特别聚焦于16x16 Patch嵌入层和Transformer编码器的工程实现。1. 理解ViT的核心设计思想传统CNN通过卷积核的滑动窗口捕捉局部特征而ViT则彻底改变了这一范式。它将图像视为由16x16像素块组成的序列每个patch被展平后作为Transformer的输入token。这种设计带来了几个关键特性全局感受野从第一层开始就能捕捉图像各区域间的长程依赖关系序列化处理图像被处理为token序列类似NLP中的单词处理方式最小化视觉先验仅保留patch划分这一最基本的空间假设在代码实现前我们需要明确几个关键参数以ViT-B/16为例参数名称值说明输入图像尺寸224x224标准ImageNet分辨率Patch大小16x16产生14x14196个patch嵌入维度768每个patch被编码为768维向量Transformer层数12编码器堆叠深度注意力头数12多头注意力机制的分头数量2. Patch嵌入层的工程实现Patch嵌入层负责将2D图像转换为1D token序列。虽然论文描述为将图像分割后线性投影但实际工程中常用卷积操作高效实现import torch import torch.nn as nn 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 # 使用卷积层实现patch分割与线性投影 self.proj nn.Conv2d( in_chans, embed_dim, kernel_sizepatch_size, stridepatch_size ) def forward(self, x): # 输入x形状: (B, C, H, W) 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这段代码的精妙之处在于使用kernel_sizestridepatch_size的卷积等效于非重叠的patch分割输出通道数embed_dim即投影维度通过flatten和transpose操作得到正确的序列形状位置编码的实践细节 ViT使用可学习的1D位置编码而非Transformer原版的固定正弦编码。这是因为图像patch的2D结构关系可以通过训练学习实际测试表明不同位置编码方式差异不大可学习参数更易于微调时调整分辨率class ViT(nn.Module): def __init__(self, img_size224, patch_size16, in_chans3, embed_dim768): super().__init__() self.patch_embed PatchEmbedding(img_size, patch_size, in_chans, embed_dim) self.num_patches self.patch_embed.num_patches # 可学习的位置编码 self.pos_embed nn.Parameter(torch.zeros(1, self.num_patches 1, embed_dim)) # 分类token self.cls_token nn.Parameter(torch.zeros(1, 1, embed_dim)) def forward(self, x): B x.shape[0] x self.patch_embed(x) # (B, N, E) cls_tokens self.cls_token.expand(B, -1, -1) x torch.cat((cls_tokens, x), dim1) # (B, 1N, E) x x self.pos_embed # 添加位置编码 return x3. Transformer编码器的模块化实现ViT的编码器与原始Transformer几乎相同但有两个微调LayerNorm位置调整到注意力块之前Pre-LN使用GELU而非ReLU作为激活函数class TransformerBlock(nn.Module): def __init__(self, embed_dim768, num_heads12, mlp_ratio4.0): super().__init__() self.norm1 nn.LayerNorm(embed_dim) self.attn nn.MultiheadAttention(embed_dim, num_heads) self.norm2 nn.LayerNorm(embed_dim) # 扩展中间层维度 hidden_dim int(embed_dim * mlp_ratio) self.mlp nn.Sequential( nn.Linear(embed_dim, hidden_dim), nn.GELU(), nn.Linear(hidden_dim, embed_dim) ) def forward(self, x): # Pre-LN结构 x_norm self.norm1(x) attn_out, _ self.attn(x_norm, x_norm, x_norm) x x attn_out x_norm self.norm2(x) mlp_out self.mlp(x_norm) x x mlp_out return x多头注意力的内存优化技巧 当处理高分辨率图像时注意力矩阵可能消耗大量内存。实用优化方法包括使用torch.nn.MultiheadAttention的优化实现混合精度训练减少显存占用梯度检查点技术gradient checkpointing4. 完整ViT模型的组装与参数分析将各组件组合成完整模型并分析其参数量与计算量class VisionTransformer(nn.Module): def __init__(self, img_size224, patch_size16, in_chans3, embed_dim768, depth12, num_heads12, num_classes1000): super().__init__() self.patch_embed PatchEmbedding(img_size, patch_size, in_chans, embed_dim) self.pos_embed nn.Parameter(torch.zeros(1, self.patch_embed.num_patches 1, embed_dim)) self.cls_token nn.Parameter(torch.zeros(1, 1, embed_dim)) # 堆叠Transformer块 self.blocks nn.Sequential(*[ TransformerBlock(embed_dim, num_heads) 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) cls_token self.cls_token.expand(x.shape[0], -1, -1) x torch.cat((cls_token, x), dim1) x x self.pos_embed x self.blocks(x) x self.norm(x) # 仅使用cls_token进行分类 cls_token_final x[:, 0] return self.head(cls_token_final)参数量计算ViT-B/16Patch嵌入层卷积核(16×16×3)×768 589,824偏置768 → 总计590,592位置编码(197×768) 151,296Transformer编码器每层注意力层QKV投影 3×(768×768) 1,769,472输出投影768×768 589,824MLP扩展768→3072→768 ≈ 2,363,904LayerNorm参数2×768×2 3,072总计每层≈4.7M分类头768×1000 768,000总参数量约86M与ResNet-152相当但计算方式完全不同。5. 训练技巧与实战建议ViT对训练条件极为敏感以下是从论文和实践中总结的关键点学习率调度使用线性warmup前5-10k步余弦退火调度基础学习率通常设为3e-4from torch.optim import AdamW optimizer AdamW(model.parameters(), lr3e-4, weight_decay0.05) # 学习率调度示例 scheduler torch.optim.lr_scheduler.CosineAnnealingLR( optimizer, T_maxtotal_steps, eta_min1e-5 )数据增强MixUp (α0.8)RandAugmentRandom Erasing标签平滑smoothing0.1分辨率调整技巧 微调时如需改变输入分辨率位置编码需进行双线性插值from torch.nn.functional import interpolate def resize_pos_embed(pos_embed, new_shape): # pos_embed: [1, N1, E] pos_embed pos_embed[:, 1:, :].transpose(1, 2) # [1, E, N] pos_embed pos_embed.view(1, -1, h, w) # 原始h,w # 双线性插值 new_pos_embed interpolate( pos_embed, sizenew_shape, modebilinear ) new_pos_embed new_pos_embed.flatten(2).transpose(1, 2) new_pos_embed torch.cat([pos_embed[:, :1, :], new_pos_embed], dim1) return new_pos_embed混合精度训练 显著减少显存占用并加速训练from torch.cuda.amp import autocast, GradScaler scaler GradScaler() for inputs, labels in dataloader: optimizer.zero_grad() with autocast(): outputs model(inputs) loss criterion(outputs, labels) scaler.scale(loss).backward() scaler.step(optimizer) scaler.update()6. 模型评估与可视化分析理解ViT工作原理的最佳方式是通过可视化注意力图可视化 提取中间层的注意力权重观察模型关注区域def visualize_attention(model, img): # 注册hook获取注意力权重 attentions [] def hook_fn(module, input, output): attentions.append(output[1]) # 注意力权重 handles [] for block in model.blocks: handles.append(block.attn.register_forward_hook(hook_fn)) with torch.no_grad(): model(img) # 移除hook for handle in handles: handle.remove() return attentionsPatch嵌入相似性 计算不同patch嵌入间的余弦相似度观察空间关系学习情况from sklearn.metrics.pairwise import cosine_similarity def patch_similarity(model, img): with torch.no_grad(): embeddings model.patch_embed(img)[0] # (N, E) sim_matrix cosine_similarity(embeddings) return sim_matrix分类token分析 通过t-SNE降维可视化不同类别cls_token的分布from sklearn.manifold import TSNE def visualize_cls_tokens(model, dataloader): cls_tokens [] labels [] for imgs, lbls in dataloader: with torch.no_grad(): features model.forward_features(imgs) # 假设有此方法 cls_tokens.append(features[:, 0]) labels.append(lbls) cls_tokens torch.cat(cls_tokens).numpy() labels torch.cat(labels).numpy() tsne TSNE(n_components2) reduced tsne.fit_transform(cls_tokens) # 绘制散点图...7. 性能优化与部署考量当需要将ViT部署到生产环境时需考虑以下优化计算图优化使用torch.jit.script或torch.jit.trace生成静态图算子融合减少内核启动开销model VisionTransformer().eval() scripted_model torch.jit.script(model) scripted_model.save(vit_scripted.pt)量化部署动态量化8bit可减少约4倍模型大小注意LayerNorm和残差连接的特殊处理quantized_model torch.quantization.quantize_dynamic( model, {nn.Linear}, dtypetorch.qint8 )ONNX导出 实现跨平台部署dummy_input torch.randn(1, 3, 224, 224) torch.onnx.export( model, dummy_input, vit.onnx, opset_version13, input_names[input], output_names[output] )内存优化策略使用梯度检查点减少训练时显存占用激活值压缩技术分布式训练策略选择from torch.utils.checkpoint import checkpoint class MemoryEfficientBlock(nn.Module): def forward(self, x): return checkpoint(self._forward, x) def _forward(self, x): # 原始forward实现 ...8. 扩展应用与前沿方向ViT的潜力远不止图像分类目标检测DETR系列将ViT与检测任务结合使用encoder-decoder架构预测边界框语义分割Segmenter纯Transformer的语义分割框架通过patch嵌入的逆操作恢复空间分辨率视频理解时空注意力机制处理视频序列将时间维度作为额外token引入自监督学习MAEMasked Autoencoder框架随机mask图像patch并预测缺失内容模型压缩DeiT通过蒸馏训练更小的ViTMobileViT轻量级混合架构# 示例简单的ViT分割头 class ViTSegmentationHead(nn.Module): def __init__(self, vit_model, num_classes): super().__init__() self.vit vit_model self.decoder nn.Sequential( nn.Conv2d(768, 256, 1), nn.Upsample(scale_factor16, modebilinear), nn.Conv2d(256, num_classes, 1) ) def forward(self, x): B, _, H, W x.shape x self.vit.patch_embed(x) # (B, N, E) x x.transpose(1, 2).view(B, -1, H//16, W//16) x self.decoder(x) return x在实现ViT的过程中最深刻的体会是其对计算资源的渴求与对大规模数据的依赖。与CNN不同ViT几乎从零开始学习所有的空间关系这既是其强大表征能力的来源也是训练难度的根源。实际应用中建议从预训练模型出发进行微调而非从头训练——除非您拥有Google级别的数据与算力。
Vision Transformer (ViT) 代码复现：PyTorch 实现 224x224 图像 16x16 Patch 嵌入
发布时间：2026/7/6 21:36:23
Vision Transformer (ViT) 代码实战从零构建224x224图像的16x16 Patch嵌入与Transformer编码器当计算机视觉领域长期被卷积神经网络CNN统治时Google Research在2020年提出的Vision TransformerViT如同一记惊雷证明了纯Transformer架构在图像识别任务中的惊人潜力。本文将带您深入ViT的核心实现细节通过PyTorch从零构建完整的ViT-B/16模型特别聚焦于16x16 Patch嵌入层和Transformer编码器的工程实现。1. 理解ViT的核心设计思想传统CNN通过卷积核的滑动窗口捕捉局部特征而ViT则彻底改变了这一范式。它将图像视为由16x16像素块组成的序列每个patch被展平后作为Transformer的输入token。这种设计带来了几个关键特性全局感受野从第一层开始就能捕捉图像各区域间的长程依赖关系序列化处理图像被处理为token序列类似NLP中的单词处理方式最小化视觉先验仅保留patch划分这一最基本的空间假设在代码实现前我们需要明确几个关键参数以ViT-B/16为例参数名称值说明输入图像尺寸224x224标准ImageNet分辨率Patch大小16x16产生14x14196个patch嵌入维度768每个patch被编码为768维向量Transformer层数12编码器堆叠深度注意力头数12多头注意力机制的分头数量2. Patch嵌入层的工程实现Patch嵌入层负责将2D图像转换为1D token序列。虽然论文描述为将图像分割后线性投影但实际工程中常用卷积操作高效实现import torch import torch.nn as nn 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 # 使用卷积层实现patch分割与线性投影 self.proj nn.Conv2d( in_chans, embed_dim, kernel_sizepatch_size, stridepatch_size ) def forward(self, x): # 输入x形状: (B, C, H, W) 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这段代码的精妙之处在于使用kernel_sizestridepatch_size的卷积等效于非重叠的patch分割输出通道数embed_dim即投影维度通过flatten和transpose操作得到正确的序列形状位置编码的实践细节 ViT使用可学习的1D位置编码而非Transformer原版的固定正弦编码。这是因为图像patch的2D结构关系可以通过训练学习实际测试表明不同位置编码方式差异不大可学习参数更易于微调时调整分辨率class ViT(nn.Module): def __init__(self, img_size224, patch_size16, in_chans3, embed_dim768): super().__init__() self.patch_embed PatchEmbedding(img_size, patch_size, in_chans, embed_dim) self.num_patches self.patch_embed.num_patches # 可学习的位置编码 self.pos_embed nn.Parameter(torch.zeros(1, self.num_patches 1, embed_dim)) # 分类token self.cls_token nn.Parameter(torch.zeros(1, 1, embed_dim)) def forward(self, x): B x.shape[0] x self.patch_embed(x) # (B, N, E) cls_tokens self.cls_token.expand(B, -1, -1) x torch.cat((cls_tokens, x), dim1) # (B, 1N, E) x x self.pos_embed # 添加位置编码 return x3. Transformer编码器的模块化实现ViT的编码器与原始Transformer几乎相同但有两个微调LayerNorm位置调整到注意力块之前Pre-LN使用GELU而非ReLU作为激活函数class TransformerBlock(nn.Module): def __init__(self, embed_dim768, num_heads12, mlp_ratio4.0): super().__init__() self.norm1 nn.LayerNorm(embed_dim) self.attn nn.MultiheadAttention(embed_dim, num_heads) self.norm2 nn.LayerNorm(embed_dim) # 扩展中间层维度 hidden_dim int(embed_dim * mlp_ratio) self.mlp nn.Sequential( nn.Linear(embed_dim, hidden_dim), nn.GELU(), nn.Linear(hidden_dim, embed_dim) ) def forward(self, x): # Pre-LN结构 x_norm self.norm1(x) attn_out, _ self.attn(x_norm, x_norm, x_norm) x x attn_out x_norm self.norm2(x) mlp_out self.mlp(x_norm) x x mlp_out return x多头注意力的内存优化技巧 当处理高分辨率图像时注意力矩阵可能消耗大量内存。实用优化方法包括使用torch.nn.MultiheadAttention的优化实现混合精度训练减少显存占用梯度检查点技术gradient checkpointing4. 完整ViT模型的组装与参数分析将各组件组合成完整模型并分析其参数量与计算量class VisionTransformer(nn.Module): def __init__(self, img_size224, patch_size16, in_chans3, embed_dim768, depth12, num_heads12, num_classes1000): super().__init__() self.patch_embed PatchEmbedding(img_size, patch_size, in_chans, embed_dim) self.pos_embed nn.Parameter(torch.zeros(1, self.patch_embed.num_patches 1, embed_dim)) self.cls_token nn.Parameter(torch.zeros(1, 1, embed_dim)) # 堆叠Transformer块 self.blocks nn.Sequential(*[ TransformerBlock(embed_dim, num_heads) 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) cls_token self.cls_token.expand(x.shape[0], -1, -1) x torch.cat((cls_token, x), dim1) x x self.pos_embed x self.blocks(x) x self.norm(x) # 仅使用cls_token进行分类 cls_token_final x[:, 0] return self.head(cls_token_final)参数量计算ViT-B/16Patch嵌入层卷积核(16×16×3)×768 589,824偏置768 → 总计590,592位置编码(197×768) 151,296Transformer编码器每层注意力层QKV投影 3×(768×768) 1,769,472输出投影768×768 589,824MLP扩展768→3072→768 ≈ 2,363,904LayerNorm参数2×768×2 3,072总计每层≈4.7M分类头768×1000 768,000总参数量约86M与ResNet-152相当但计算方式完全不同。5. 训练技巧与实战建议ViT对训练条件极为敏感以下是从论文和实践中总结的关键点学习率调度使用线性warmup前5-10k步余弦退火调度基础学习率通常设为3e-4from torch.optim import AdamW optimizer AdamW(model.parameters(), lr3e-4, weight_decay0.05) # 学习率调度示例 scheduler torch.optim.lr_scheduler.CosineAnnealingLR( optimizer, T_maxtotal_steps, eta_min1e-5 )数据增强MixUp (α0.8)RandAugmentRandom Erasing标签平滑smoothing0.1分辨率调整技巧 微调时如需改变输入分辨率位置编码需进行双线性插值from torch.nn.functional import interpolate def resize_pos_embed(pos_embed, new_shape): # pos_embed: [1, N1, E] pos_embed pos_embed[:, 1:, :].transpose(1, 2) # [1, E, N] pos_embed pos_embed.view(1, -1, h, w) # 原始h,w # 双线性插值 new_pos_embed interpolate( pos_embed, sizenew_shape, modebilinear ) new_pos_embed new_pos_embed.flatten(2).transpose(1, 2) new_pos_embed torch.cat([pos_embed[:, :1, :], new_pos_embed], dim1) return new_pos_embed混合精度训练 显著减少显存占用并加速训练from torch.cuda.amp import autocast, GradScaler scaler GradScaler() for inputs, labels in dataloader: optimizer.zero_grad() with autocast(): outputs model(inputs) loss criterion(outputs, labels) scaler.scale(loss).backward() scaler.step(optimizer) scaler.update()6. 模型评估与可视化分析理解ViT工作原理的最佳方式是通过可视化注意力图可视化 提取中间层的注意力权重观察模型关注区域def visualize_attention(model, img): # 注册hook获取注意力权重 attentions [] def hook_fn(module, input, output): attentions.append(output[1]) # 注意力权重 handles [] for block in model.blocks: handles.append(block.attn.register_forward_hook(hook_fn)) with torch.no_grad(): model(img) # 移除hook for handle in handles: handle.remove() return attentionsPatch嵌入相似性 计算不同patch嵌入间的余弦相似度观察空间关系学习情况from sklearn.metrics.pairwise import cosine_similarity def patch_similarity(model, img): with torch.no_grad(): embeddings model.patch_embed(img)[0] # (N, E) sim_matrix cosine_similarity(embeddings) return sim_matrix分类token分析 通过t-SNE降维可视化不同类别cls_token的分布from sklearn.manifold import TSNE def visualize_cls_tokens(model, dataloader): cls_tokens [] labels [] for imgs, lbls in dataloader: with torch.no_grad(): features model.forward_features(imgs) # 假设有此方法 cls_tokens.append(features[:, 0]) labels.append(lbls) cls_tokens torch.cat(cls_tokens).numpy() labels torch.cat(labels).numpy() tsne TSNE(n_components2) reduced tsne.fit_transform(cls_tokens) # 绘制散点图...7. 性能优化与部署考量当需要将ViT部署到生产环境时需考虑以下优化计算图优化使用torch.jit.script或torch.jit.trace生成静态图算子融合减少内核启动开销model VisionTransformer().eval() scripted_model torch.jit.script(model) scripted_model.save(vit_scripted.pt)量化部署动态量化8bit可减少约4倍模型大小注意LayerNorm和残差连接的特殊处理quantized_model torch.quantization.quantize_dynamic( model, {nn.Linear}, dtypetorch.qint8 )ONNX导出 实现跨平台部署dummy_input torch.randn(1, 3, 224, 224) torch.onnx.export( model, dummy_input, vit.onnx, opset_version13, input_names[input], output_names[output] )内存优化策略使用梯度检查点减少训练时显存占用激活值压缩技术分布式训练策略选择from torch.utils.checkpoint import checkpoint class MemoryEfficientBlock(nn.Module): def forward(self, x): return checkpoint(self._forward, x) def _forward(self, x): # 原始forward实现 ...8. 扩展应用与前沿方向ViT的潜力远不止图像分类目标检测DETR系列将ViT与检测任务结合使用encoder-decoder架构预测边界框语义分割Segmenter纯Transformer的语义分割框架通过patch嵌入的逆操作恢复空间分辨率视频理解时空注意力机制处理视频序列将时间维度作为额外token引入自监督学习MAEMasked Autoencoder框架随机mask图像patch并预测缺失内容模型压缩DeiT通过蒸馏训练更小的ViTMobileViT轻量级混合架构# 示例简单的ViT分割头 class ViTSegmentationHead(nn.Module): def __init__(self, vit_model, num_classes): super().__init__() self.vit vit_model self.decoder nn.Sequential( nn.Conv2d(768, 256, 1), nn.Upsample(scale_factor16, modebilinear), nn.Conv2d(256, num_classes, 1) ) def forward(self, x): B, _, H, W x.shape x self.vit.patch_embed(x) # (B, N, E) x x.transpose(1, 2).view(B, -1, H//16, W//16) x self.decoder(x) return x在实现ViT的过程中最深刻的体会是其对计算资源的渴求与对大规模数据的依赖。与CNN不同ViT几乎从零开始学习所有的空间关系这既是其强大表征能力的来源也是训练难度的根源。实际应用中建议从预训练模型出发进行微调而非从头训练——除非您拥有Google级别的数据与算力。
