看 ViT 源码之前最难的往往不是 Self-Attention 公式而是一个更基础的问题图像是二维的Transformer 是处理一维序列的这两者之间到底怎么“翻译”很多读者从 NLP 的 Transformer 教程入手理解词向量、理解 QKV 都很顺利。但一接触到 Vision TransformerViT立刻卡在 Patch Embedding 上。原因很简单从 CNN 的“卷积滑窗”思维切到“图像分块拉直”思维中间缺少一座桥。这座桥就是本文要彻底拆开的 Patch Embedding 和 Forward 流程。本文的核心判断是理解了 Patch Embedding就理解了 ViT 的一半。剩下的一半才是你熟悉的 Transformer Encoder。我会用 PyTorch 从零手撕 ViT 的完整代码把输入图像从[B, 3, 224, 224]变成[B, 197, 192]的每一步 shape 变化都讲清楚最后给出可运行的完整示例和常见坑位排查。读完你可以自己复现一个 ImageNet 分类用的 ViT 最小实现。1. 为什么说 ViT 的核心不是 Attention而是 Patch Embedding先抛一个反直觉的观点ViT 最关键的创新不是把 NLP 的 Transformer 搬到图像上而是设计了一个足够简单的“图像转序列”方案让 Transformer 不需要依赖 CNN 也能直接处理图像。在 ViT 之前图像领域是 CNN 的天下。CNN 靠卷积核在空间上滑动天然具有局部性和平移等变性这是它擅长处理图像的原因。但 CNN 也有明显局限卷积核感受野有限想捕捉长距离依赖必须堆很多层。而 Transformer 的自注意力机制第一步就能让任意两个位置直接交互天然适合建模全局依赖。那为什么不直接把图像的每个像素当成一个“词”送给 Transformer一个 224×224 的 RGB 图像有 150528 个像素值如果每个像素是一个 token序列长度就超过 15 万。自注意力的复杂度是 O(n²)这个长度直接会让显存爆炸而且在像素级别建模全局关系既浪费计算量也没有必要。ViT 的答案是把图像切成固定大小的 Patch块每个 Patch 看作一个“视觉词”。一张 224×224 的图像切成 16×16 的 Patch就得到 14×14196 个 Patch。序列长度从超过 15 万降到 196这是质的变化。这个过程叫 Patch Embedding。它同时完成了三件事切块把图像切成不重叠的 Patch。展平把每个 Patch 的像素展平成一维向量。线性映射把展平后的向量投影到 Transformer 需要的维度空间。用一句话类比ViT 把读图变成读文章。每个 Patch 相当于一个词Patch 的空间排列顺序相当于词的顺序Transformer 就像一个阅读理解模型读完整张图后输出分类结果。理解了这一点就理解了为什么 Patch Embedding 是 ViT 的核心。2. Transformer 基础Self-Attention 到底在算什么在写 ViT 代码之前必须先弄清 Transformer Encoder 每一层在做什么。否则即使代码跑通了你也不知道它在算什么。2.1 QKV 是什么自注意力机制的输入是一个序列每个 token 对应一个特征向量。在 ViT 里这个特征向量就是一个 Patch 经过 Patch Embedding 后的向量。对每个特征向量 x模型通过三个可学习的线性变换分别生成 Query、Key、ValueQuery查询表示“我在找什么”。Key键表示“我有什么特征可以被匹配”。Value值表示“我实际携带的信息”。Query 和 Key 做点积得到注意力分数决定“我该关注谁”。再用这个分数对 Value 做加权求和得到当前 token 融合全局信息后的新表示。2.2 Attention 公式注意力分数的计算方式如下Attention(Q, K, V) softmax(Q · K^T / √d_k) · V其中 d_k 是每个注意力头的维度。除以 √d_k 的作用是防止点积结果过大导致 softmax 梯度消失。直观理解就是当向量维度很大时点积结果的方差会变大除以 √d_k 可以把方差拉回 1 附近让 softmax 的输入在一个合理区间。Self-Attention自注意力的意思是Q、K、V 都来自同一输入序列。每个 token 在计算输出时会参考序列里所有其他 token 的信息。这就是 Transformer 能建模全局依赖的根本原因。2.3 多头注意力的作用多头注意力Multi-Head Attention是多个自注意力头的组合。每个头有独立的 QKV 投影可以从不同子空间学习不同的注意力模式。例如在图像里一个头可能关注“这个 Patch 和左边 Patch 的颜色关系”另一个头可能关注“这个 Patch 和远处 Patch 的形状关系”。代码里实现多头注意力的常用方式是把特征维度切成 num_heads 份每个头独立计算注意力最后拼接在一起再过一次线性层。3. ViT 的整体架构图像是怎么变成 token 的ViT 的整体流程可以拆成 5 个阶段原始图像 → Patch Embedding → 拼接 CLS Token → 加入位置编码 → Transformer Encoder → 分类输出3.1 Patch Embedding输入是[B, 3, H, W]的图像输出是[B, N, embed_dim]。其中 N 是 Patch 数量等于(H / patch_size) × (W / patch_size)。embed_dim 是 Transformer 的特征维度也是每个 Patch 经过线性映射后的向量长度。3.2 CLS TokenViT 参考 BERT 的做法在序列开头拼接一个特殊的 CLS Token。这个 token 不来自任何 Patch它的初始值是可学习的参数和 Patch 一起参与 Transformer 前向计算。经过多层 Transformer 编码后CLS Token 的最后一层输出可以被理解为“整张图像的全局特征表示”分类头就接在这个位置。为什么不直接对所有 Patch 的输出取平均效果上两者都可能工作但 CLS Token 的设计有两个实际好处分类头的输入位置固定不需要额外写聚合逻辑。CLS Token 通过注意力机制主动聚合全图信息而不是简单平均。3.3 Position Embedding位置编码Transformer 的 self-attention 本身不感知 token 的先后顺序。无论把第 1 个 Patch 和第 196 个 Patch 调换位置注意力结果都是一样的。但图像的 Patch 顺序显然携带空间语义信息所以必须把位置信息注入到输入序列中。ViT 采用可学习的位置编码learnable position embedding它是一个 shape 为[1, N1, embed_dim]的参数矩阵直接加到输入序列上。注意这里 N1 是因为多了一个 CLS Token。3.4 Transformer Encoder这是 ViT 的特征提取主体由多个相同的 Block 堆叠而成。每个 Block 包含LayerNorm Multi-Head Attention 残差连接LayerNorm MLP 残差连接3.5 分类头取 Encoder 输出序列中 CLS Token 位置的特征向量过一层线性层得到分类 logits。4. 环境准备与最小依赖在动手写代码之前先确认环境。本文的代码基于深度学习框架 PyTorch不依赖额外的第三方库。# 建议使用 Python 3.8 及以上版本 # 安装 PyTorch版本请以官方安装命令为准 pip install torch torchvision验证安装是否成功import torch print(torch.__version__) print(torch.cuda.is_available())需要注意两点本文的代码在 CPU 上也能运行。显存紧张的情况下建议先用小配置跑通逻辑再放大模型规模。如果环境确实没有 PyTorch请先前往 PyTorch 官网根据操作系统、Python 版本和 CUDA 版本生成对应的安装命令不要使用来源不明的安装源。5. Patch Embedding 完整代码实现Patch Embedding 在代码层面有两种实现方式一种是直观的“手动切块 线性映射”另一种是工程上更高效的“Conv2d 一步实现”。两种方式数学上完全等价我建议都掌握。5.1 手动切块 线性映射最直观的实现import torch import torch.nn as nn class PatchEmbeddingNaive(nn.Module): 最直观的 Patch Embedding 实现 1. 把图像切成不重叠的 patch 2. 每个 patch 展平成向量 3. 过线性层映射到 embed_dim def __init__(self, in_channels3, patch_size16, embed_dim768): super().__init__() self.patch_size patch_size # patch 展平后的维度 通道数 × patch 高 × patch 宽 flatten_dim in_channels * patch_size * patch_size self.linear nn.Linear(flatten_dim, embed_dim) def forward(self, x): B, C, H, W x.shape p self.patch_size assert H % p 0 and W % p 0, \ f输入尺寸 {H}x{W} 必须能被 patch_size{p} 整除 # unfold 可以看作“滑动窗口取块”按步长 p 切出 p×p 的块 # 结果 shape: [B, C, H/p, W/p, p, p] x x.unfold(2, p, p).unfold(3, p, p) # 调整维度顺序把空间分块放在前两维 # [B, C, H/p, W/p, p, p] - [B, H/p, W/p, C, p, p] x x.permute(0, 2, 3, 1, 4, 5).contiguous() # 把 C×p×p 展平成一个向量 # [B, H/p, W/p, C*p*p] x x.flatten(3) # 把 H/p 和 W/p 合并成 Patch 数量 N # [B, N, C*p*p] 其中 N H/p × W/p x x.flatten(1, 2) # 线性映射到 embed_dim # [B, N, embed_dim] x self.linear(x) return x手动实现的核心是理解 unfold 的含义。unfold(2, p, p)表示在第 2 维高度方向上以窗口大小为 p、步长为 p 取出块。第 3 维同理。这样一次操作就把整张图切成了 (H/p) × (W/p) 个 Patch。5.2 用 Conv2d 实现工程上最优雅的写法class PatchEmbedding(nn.Module): 工程上最推荐的 Patch Embedding 实现。 用 kernel_sizepatch_size, stridepatch_size 的卷积一步完成 “切块 展平 线性映射”。 def __init__(self, in_channels3, patch_size16, embed_dim768): super().__init__() self.patch_size patch_size # 这一步卷积同时完成了三件事 # 1. kernel_sizepatch_size 相当于取一个 patch 大小的窗口 # 2. stridepatch_size 相当于 patch 之间不重叠 # 3. 输出通道数为 embed_dim相当于每个 patch 线性映射到 embed_dim self.proj nn.Conv2d( in_channels, embed_dim, kernel_sizepatch_size, stridepatch_size ) def forward(self, x): B, C, H, W x.shape assert H % self.patch_size 0 and W % self.patch_size 0, \ f输入尺寸 {H}x{W} 必须能被 patch_size{self.patch_size} 整除 # 卷积输出: [B, embed_dim, H/p, W/p] x self.proj(x) # 展平空间维度: [B, embed_dim, N]N H/p × W/p x x.flatten(2) # 转成序列格式: [B, N, embed_dim] x x.transpose(1, 2) return x为什么 Conv2d 能做到和手动切块等价核心在卷积核的滑动方式。当 kernel_size patch_size、stride patch_size 时卷积核每次正好覆盖一个 Patch且 Patch 之间没有重叠。每个输出位置的值就是该 Patch 内的像素值和卷积核做内积的结果。多个输出通道相当于多个线性映射等价于nn.Linear(flatten_dim, embed_dim)。从性能角度看Conv2d 底层是高度优化的矩阵运算比手动 unfold 之后再过 Linear 更快。因此在实际工程代码如 TIMM、HuggingFace Transformers中基本都是用 Conv2d 实现 Patch Embedding。5.3 两种实现方式的对比对比维度手动切块 LinearConv2d 实现可读性容易理解适合教学简洁高效适合工程计算速度相对较慢更快底层高度优化扩展性改逻辑更直观改 patch_size 只需改参数实际项目较少使用TIMM 等库的标准做法我的建议是入门阶段用手动实现理解原理实际写项目时用 Conv2d 版本。6. 从零手撕 ViT完整 Forward 代码这一节是核心。我会按照“Attention → Block → Encoder → 完整 ViT”的顺序从零搭建 Vision Transformer。6.1 多头自注意力模块class Attention(nn.Module): 多头自注意力模块。 输入输出 shape 相同: [B, N, dim] def __init__(self, dim, num_heads8, qkv_biasFalse, attn_drop0.0): super().__init__() self.num_heads num_heads self.scale (dim // num_heads) ** -0.5 # 一个线性层同时生成 Q、K、V输出维度是 3 * dim self.qkv nn.Linear(dim, dim * 3, biasqkv_bias) self.attn_drop nn.Dropout(attn_drop) self.proj nn.Linear(dim, dim) def forward(self, x): B, N, C x.shape # 生成 QKV并拆成多头 # 先 reshape 成 [B, N, 3, num_heads, C // num_heads] qkv self.qkv(x).reshape(B, N, 3, self.num_heads, C // self.num_heads) # 再 permute 成 [3, B, num_heads, N, C // num_heads] qkv qkv.permute(2, 0, 3, 1, 4) q, k, v qkv[0], qkv[1], qkv[2] # 计算注意力分数: [B, num_heads, N, N] attn (q k.transpose(-2, -1)) * self.scale attn attn.softmax(dim-1) attn self.attn_drop(attn) # 注意力加权求和: [B, num_heads, N, C // num_heads] x attn v # 合并多头: [B, N, C] x x.transpose(1, 2).reshape(B, N, C) x self.proj(x) return x这段代码最需要理解的是维度变换。reshape把 C 维拆成3, num_heads, C // num_heads三部分permute把 QKV 这个维度挪到最前面方便用qkv[0]、qkv[1]、qkv[2]分别取出 Q、K、V。6.2 MLP 与 Transformer Blockclass Mlp(nn.Module): Transformer Block 中的前馈网络。 通常包含两个线性层和一个 GELU 激活。 def __init__(self, in_features, hidden_featuresNone, dropout0.0): super().__init__() hidden_features hidden_features or in_features self.fc1 nn.Linear(in_features, hidden_features) self.act nn.GELU() self.fc2 nn.Linear(hidden_features, in_features) def forward(self, x): x self.fc1(x) x self.act(x) x self.fc2(x) return x class Block(nn.Module): 标准 Transformer Encoder Block: LayerNorm - Attention - 残差连接 - LayerNorm - MLP - 残差连接 def __init__(self, dim, num_heads, mlp_ratio4.0): super().__init__() self.norm1 nn.LayerNorm(dim) self.attn Attention(dim, num_headsnum_heads) self.norm2 nn.LayerNorm(dim) self.mlp Mlp(in_featuresdim, hidden_featuresint(dim * mlp_ratio)) def forward(self, x): # Pre-LN 结构先 LayerNorm 再 Attention残差连接在外层 x x self.attn(self.norm1(x)) x x self.mlp(self.norm2(x)) return x注意这里使用的是 Pre-LN 结构即 LayerNorm 放在 Attention 和 MLP 之前。研究表明Pre-LN 比 Post-LN 更稳定训练时不容易出现梯度爆炸这也是 ViT 论文和 TIMM 库常用做法。Block 里两次残差连接的 shape 完全一致都是[B, N, dim]这是 Transformer 能堆叠几十层的基础。6.3 完整 ViT 模型class VisionTransformer(nn.Module): 完整 Vision Transformer: Patch Embedding - CLS Token - Position Embedding - Encoder - 分类头 def __init__(self, img_size224, patch_size16, in_channels3, num_classes1000, embed_dim768, depth12, num_heads12, mlp_ratio4.0): super().__init__() self.patch_embed PatchEmbedding(in_channels, patch_size, embed_dim) # 计算 Patch 数量例如 224/16 14N 14 × 14 196 num_patches (img_size // patch_size) ** 2 # 可学习的 CLS Tokenshape: [1, 1, embed_dim] self.cls_token nn.Parameter(torch.zeros(1, 1, embed_dim)) # 可学习的位置编码包含 CLS Token 的位置所以需要 1 # shape: [1, num_patches 1, embed_dim] self.pos_embed nn.Parameter(torch.zeros(1, num_patches 1, embed_dim)) self.pos_drop nn.Dropout(0.0) # 堆叠 Transformer Encoder Block self.blocks nn.ModuleList([ Block(embed_dim, num_heads, mlp_ratio) for _ in range(depth) ]) self.norm nn.LayerNorm(embed_dim) # 分类头默认接 ImageNet 1000 类 self.head nn.Linear(embed_dim, num_classes) def forward_features(self, x): B x.shape[0] # 1. Patch Embedding: [B, 3, 224, 224] - [B, 196, 768] x self.patch_embed(x) # 2. 拼接 CLS Token # cls_token: [1, 1, 768] - expand 到 [B, 1, 768] cls_token self.cls_token.expand(B, -1, -1) x torch.cat([cls_token, x], dim1) # [B, 197, 768] # 3. 加入位置编码 x x self.pos_embed # [B, 197, 768] # 4. 过 Transformer Encoder for blk in self.blocks: x blk(x) # shape 始终不变 x self.norm(x) return x def forward(self, x): # 提取特征 x self.forward_features(x) # [B, 197, 768] # 取 CLS Token 的输出特征 x x[:, 0] # [B, 768] # 分类头 x self.head(x) # [B, num_classes] return x6.4 Forward 全流程 shape 变化用一个小型配置跑一次完整前向方便观察每一步 shape# 小型配置方便在 CPU 上跑通流程 model VisionTransformer( img_size224, patch_size16, num_classes10, embed_dim192, depth4, num_heads4, mlp_ratio4.0 ) # 模拟输入2 张 224×224 RGB 图像 x torch.randn(2, 3, 224, 224) # 前向传播 out model(x) print(输出 shape:, out.shape) # torch.Size([2, 10])对应每一阶段的 shape 变化如下阶段Shape说明输入图像[2, 3, 224, 224]B23 通道224×224Patch Embedding 后[2, 196, 192]N196embed_dim192拼接 CLS Token 后[2, 197, 192]序列长度 1加入位置编码后[2, 197, 192]与 pos_embed 相加Encoder 每一层后[2, 197, 192]共 4 层shape 不变取 CLS 特征[2, 192]取第 0 个位置分类输出[2, 10]过 Linear 分类头这个表是整个 ViT 前向过程的骨架。只要你的代码 shape 变化和表一致说明模型结构搭建正确。7. 运行结果与效果验证光看 shape 还不够最好能打印出每一层输出的实际值帮助确认模型是否真的在工作。写一个调试版子类在 forward 过程中插入打印语句class ViTDebug(VisionTransformer): def forward_features(self, x): B x.shape[0] print(f[0] 输入图像: {x.shape}) x self.patch_embed(x) print(f[1] Patch Embedding: {x.shape}) cls_token self.cls_token.expand(B, -1, -1) x torch.cat([cls_token, x], dim1) print(f[2] 拼接 CLS Token: {x.shape}, CLS 值: {x[:, 0, :5]}) x x self.pos_embed print(f[3] 加入位置编码: {x.shape}) for i, blk in enumerate(self.blocks): x blk(x) print(f[4] Block {i1}: {x.shape}) x self.norm(x) print(f[5] 最终 LayerNorm: {x.shape}) return x model_debug ViTDebug( img_size224, patch_size16, num_classes10, embed_dim192, depth4, num_heads4 ) x torch.randn(2, 3, 224, 224) out model_debug(x) print(分类输出:, out.shape)运行后应该看到类似输出[0] 输入图像: torch.Size([2, 3, 224, 224]) [1] Patch Embedding: torch.Size([2, 196, 192]) [2] 拼接 CLS Token: torch.Size([2, 197, 192]) [3] 加入位置编码: torch.Size([2, 197, 192]) [4] Block 1: torch.Size([2, 197, 192]) [4] Block 2: torch.Size([2, 197, 192]) [4] Block 3: torch.Size([2, 197, 192]) [4] Block 4: torch.Size([2, 197, 192]) [5] 最终 LayerNorm: torch.Size([2, 197, 192]) 分类输出: torch.Size([2, 10])判断模型是否搭建成功的标准所有 shape 和上面表格一致。前向传播没有报错。分类输出 logits 不是全零或 NaN。CLS Token 的值在整个过程中持续变化说明它确实在参与注意力计算。如果运行失败先不要急着改模型结构。按下面第 8 节排查。我还可以用一个快速实验验证模型“真的能学”用随机生成的标签做一步反向传播看 loss 是否能下降。这个能确认模型参数确实在更新。import torch.nn.functional as F # 随机生成一对数据和标签仅用于验证梯度流 x torch.randn(4, 3, 224, 224) y torch.randint(0, 10, (4,)) model VisionTransformer( img_size224, patch_size16, num_classes10, embed_dim192, depth4, num_heads4 ) optimizer torch.optim.AdamW(model.parameters(), lr3e-4) loss F.cross_entropy(model(x), y) print(初始 loss:, loss.item()) optimizer.zero_grad() loss.backward() optimizer.step() loss2 F.cross_entropy(model(x), y) print(更新后 loss:, loss2.item())如果 loss 能正常下降说明整个模型从 forward 到 backward 的计算图是完全通的。这一步也是后续接入真实训练流程的起点。8. 常见问题与排查思路手撕 ViT 的过程中有几个问题几乎每个人都会遇到。我把常见问题整理成表方便对照排查。问题现象可能原因排查方式解决方案输入图像尺寸不能整除 patch_sizePatch Embedding 的 assert 拦截查看报错信息中的 H、W、patch_size调整 img_size 或 patch_size或者对输入图像做 resize拼接 CLS Token 后维度报错cls_token 的 embed_dim 与 patch_embed 输出维度不一致打印 patch_embed 输出和 cls_token shape确保 embed_dim 参数在所有模块中保持一致位置编码相加报错pos_embed 的序列长度不是 num_patches1打印 x 和 pos_embed 的 shape检查 num_patches 计算是否正确记得 1注意力分数为 NaN学习率过大或初始化不当打印 attn 的统计值降低学习率检查数据是否需要归一化显存不足模型配置过大、batch_size 过大、序列太长查看 GPU 显存占用减小 embed_dim、depth 或 batch_size训练 loss 不下降没有使用 warmup、学习率不合适、数据没归一化打印梯度范数使用 AdamW warmup学习率设为 3e-4 量级分类准确率远低于 CNN数据量不足、训练轮次不够、缺少数据增强观察训练集和验证集 loss 差距增加训练数据、使用随机裁剪/翻转/MixUp 等增强策略下面挑几个重点问题展开说明。8.1 为什么维度老是差 1很多新手第一次跑通模型后发现 pos_embed 相加时报错原因是“我明明算了 num_patches为什么加不上去”问题在于忘了 CLS Token 也占一个位置。Patch Embedding 输出 N 个 Patch拼接 CLS 后序列长度变成 N1所以位置编码的序列长度也必须是 N1。这个 “1” 是 ViT 里最容易漏掉的地方。8.2 为什么用 Conv2d 当 Patch Embedding 不感觉奇怪吗这是初学者最常问的问题之一。Conv2d 通常是用来提取特征的为什么在这里变成了“切块”工具关键在 kernel_size 和 stride 都等于 patch_size。这会导致卷积核覆盖的区域互不重叠输出的每一个空间位置恰好对应原图的一个 Patch。它做的本质上是一个“不重叠的局部线性映射”和先展平再过线性层完全等价。理解这一点后你会发现 Conv2d 其实是一个精巧的工程实现而不是一个“玄学”操作。8.3 为什么我的 ViT 训练不收敛如果在随机数据上都看不到 loss 明显下降先检查三点数据有没有归一化到 [0,1] 或标准化ViT 对输入尺度比较敏感。学习率是否合理ViT 一般用 AdamW初始学习率在 3e-4 左右并配合 warmup。类别数是否正确如果数据集只有 10 类但分类头输出 1000 类loss 很难下降。8.4 显存不够怎么办ViT 的显存占用明显高于同规模 CNN尤其是自注意力矩阵[B, num_heads, N, N]会占用大量显存。解决办法减小 patch_size 会增大序列长度 N显存会显著上升增大 patch_size 则相反。减小 embed_dim 和 depth 是最直接的降显存手段。减少 batch_size或者使用梯度累积。8.5 为什么 Transformer 在图像上的收敛速度比 CNN 慢ViT 没有 CNN 那种强归纳偏置局部性、平移等变性需要更多数据才能学到同样的规律。在小型数据集上ViT 通常打不过同规模 ResNet。这也是为什么后来出现了 DeiT蒸馏 数据增强、Swin Transformer窗口自注意力 层级结构等改进方案。9. 工程建议与后续学习方向9.1 位置编码不是只有一种选择本文实现的是 ViT 原论文的 1D 可学习位置编码。实际工程中还有几种选择2D 位置编码针对图像的行列坐标分别建模但 ViT 论文实验显示 1D 效果已经够好因为 Transformer 可以通过注意力自己学到空间结构。相对位置编码在注意力分数计算中引入 Patch 之间的相对位置偏移常用于 Swin Transformer 等改进模型。正弦位置编码NLP Transformer 的经典方案不需要训练但图像场景中较少使用。如果要做下游目标检测或分割任务位置编码的选择会产生明显影响。建议先跑通 1D 可学习版本再根据任务需求替换。9.2 优化器与训练策略ViT 训练有两条关键经验优化器推荐 AdamW基础学习率通常设为 3e-4 量级batch size 较大时可适当提高。建议使用 warmup即学习率先从小值线性增长到目标值再按余弦退火衰减。这主要是因为 Transformer 对早期训练不稳定比较敏感。如果显存有限batch size 很小可以考虑使用梯度累积保证等效 batch size 足够大。9.3 数据量与正则化ViT 在小数据集上容易过拟合。如果在 CIFAR-10 这类小数据集上训练建议使用随机裁剪、随机翻转、MixUp、CutMix 等数据增强。增加 Dropout 和 DropPath。预训练后微调这是最有效的方案先在 ImageNet-21k 等大数据集上预训练再在下游任务上做 fine-tune。9.4 从 ViT 走向更复杂的视觉 Transformer跑通本文的 ViT 之后下一步建议按这个顺序深入DeiT用知识蒸馏减少 ViT 对超大数据集的依赖训练策略有参考价值。Swin Transformer引入窗口注意力和层级设计解决 ViT 计算复杂度高、缺乏多尺度特征的问题。MAE用自监督掩码重建方式预训练 ViT展示了 Transformer 在图像上的强大表征能力。CLIPViT 作为视觉编码器和文本编码器对齐统一了图文表示。DINO自监督 ViT得到的特征具有很好的语义分割和物体发现能力。理解 ViT 之后这些模型的很多代码你都能快速看懂因为它们都是“Patch Embedding Transformer Encoder”这个基本骨架的变体。建议你把这篇文章里的完整代码自己敲一遍而不是直接复制。敲的过程中在纸上画出每一步的 shape 变化尤其是 QKV 拆分和注意力矩阵的那几个维度。当你做到不用看代码也能说出[B, num_heads, N, N]的来龙去脉时Transformer 的核心原理就真正属于你了。 从零手撕ViT：Patch Embedding与完整PyTorch实现
看 ViT 源码之前最难的往往不是 Self-Attention 公式而是一个更基础的问题图像是二维的Transformer 是处理一维序列的这两者之间到底怎么“翻译”很多读者从 NLP 的 Transformer 教程入手理解词向量、理解 QKV 都很顺利。但一接触到 Vision TransformerViT立刻卡在 Patch Embedding 上。原因很简单从 CNN 的“卷积滑窗”思维切到“图像分块拉直”思维中间缺少一座桥。这座桥就是本文要彻底拆开的 Patch Embedding 和 Forward 流程。本文的核心判断是理解了 Patch Embedding就理解了 ViT 的一半。剩下的一半才是你熟悉的 Transformer Encoder。我会用 PyTorch 从零手撕 ViT 的完整代码把输入图像从[B, 3, 224, 224]变成[B, 197, 192]的每一步 shape 变化都讲清楚最后给出可运行的完整示例和常见坑位排查。读完你可以自己复现一个 ImageNet 分类用的 ViT 最小实现。1. 为什么说 ViT 的核心不是 Attention而是 Patch Embedding先抛一个反直觉的观点ViT 最关键的创新不是把 NLP 的 Transformer 搬到图像上而是设计了一个足够简单的“图像转序列”方案让 Transformer 不需要依赖 CNN 也能直接处理图像。在 ViT 之前图像领域是 CNN 的天下。CNN 靠卷积核在空间上滑动天然具有局部性和平移等变性这是它擅长处理图像的原因。但 CNN 也有明显局限卷积核感受野有限想捕捉长距离依赖必须堆很多层。而 Transformer 的自注意力机制第一步就能让任意两个位置直接交互天然适合建模全局依赖。那为什么不直接把图像的每个像素当成一个“词”送给 Transformer一个 224×224 的 RGB 图像有 150528 个像素值如果每个像素是一个 token序列长度就超过 15 万。自注意力的复杂度是 O(n²)这个长度直接会让显存爆炸而且在像素级别建模全局关系既浪费计算量也没有必要。ViT 的答案是把图像切成固定大小的 Patch块每个 Patch 看作一个“视觉词”。一张 224×224 的图像切成 16×16 的 Patch就得到 14×14196 个 Patch。序列长度从超过 15 万降到 196这是质的变化。这个过程叫 Patch Embedding。它同时完成了三件事切块把图像切成不重叠的 Patch。展平把每个 Patch 的像素展平成一维向量。线性映射把展平后的向量投影到 Transformer 需要的维度空间。用一句话类比ViT 把读图变成读文章。每个 Patch 相当于一个词Patch 的空间排列顺序相当于词的顺序Transformer 就像一个阅读理解模型读完整张图后输出分类结果。理解了这一点就理解了为什么 Patch Embedding 是 ViT 的核心。2. Transformer 基础Self-Attention 到底在算什么在写 ViT 代码之前必须先弄清 Transformer Encoder 每一层在做什么。否则即使代码跑通了你也不知道它在算什么。2.1 QKV 是什么自注意力机制的输入是一个序列每个 token 对应一个特征向量。在 ViT 里这个特征向量就是一个 Patch 经过 Patch Embedding 后的向量。对每个特征向量 x模型通过三个可学习的线性变换分别生成 Query、Key、ValueQuery查询表示“我在找什么”。Key键表示“我有什么特征可以被匹配”。Value值表示“我实际携带的信息”。Query 和 Key 做点积得到注意力分数决定“我该关注谁”。再用这个分数对 Value 做加权求和得到当前 token 融合全局信息后的新表示。2.2 Attention 公式注意力分数的计算方式如下Attention(Q, K, V) softmax(Q · K^T / √d_k) · V其中 d_k 是每个注意力头的维度。除以 √d_k 的作用是防止点积结果过大导致 softmax 梯度消失。直观理解就是当向量维度很大时点积结果的方差会变大除以 √d_k 可以把方差拉回 1 附近让 softmax 的输入在一个合理区间。Self-Attention自注意力的意思是Q、K、V 都来自同一输入序列。每个 token 在计算输出时会参考序列里所有其他 token 的信息。这就是 Transformer 能建模全局依赖的根本原因。2.3 多头注意力的作用多头注意力Multi-Head Attention是多个自注意力头的组合。每个头有独立的 QKV 投影可以从不同子空间学习不同的注意力模式。例如在图像里一个头可能关注“这个 Patch 和左边 Patch 的颜色关系”另一个头可能关注“这个 Patch 和远处 Patch 的形状关系”。代码里实现多头注意力的常用方式是把特征维度切成 num_heads 份每个头独立计算注意力最后拼接在一起再过一次线性层。3. ViT 的整体架构图像是怎么变成 token 的ViT 的整体流程可以拆成 5 个阶段原始图像 → Patch Embedding → 拼接 CLS Token → 加入位置编码 → Transformer Encoder → 分类输出3.1 Patch Embedding输入是[B, 3, H, W]的图像输出是[B, N, embed_dim]。其中 N 是 Patch 数量等于(H / patch_size) × (W / patch_size)。embed_dim 是 Transformer 的特征维度也是每个 Patch 经过线性映射后的向量长度。3.2 CLS TokenViT 参考 BERT 的做法在序列开头拼接一个特殊的 CLS Token。这个 token 不来自任何 Patch它的初始值是可学习的参数和 Patch 一起参与 Transformer 前向计算。经过多层 Transformer 编码后CLS Token 的最后一层输出可以被理解为“整张图像的全局特征表示”分类头就接在这个位置。为什么不直接对所有 Patch 的输出取平均效果上两者都可能工作但 CLS Token 的设计有两个实际好处分类头的输入位置固定不需要额外写聚合逻辑。CLS Token 通过注意力机制主动聚合全图信息而不是简单平均。3.3 Position Embedding位置编码Transformer 的 self-attention 本身不感知 token 的先后顺序。无论把第 1 个 Patch 和第 196 个 Patch 调换位置注意力结果都是一样的。但图像的 Patch 顺序显然携带空间语义信息所以必须把位置信息注入到输入序列中。ViT 采用可学习的位置编码learnable position embedding它是一个 shape 为[1, N1, embed_dim]的参数矩阵直接加到输入序列上。注意这里 N1 是因为多了一个 CLS Token。3.4 Transformer Encoder这是 ViT 的特征提取主体由多个相同的 Block 堆叠而成。每个 Block 包含LayerNorm Multi-Head Attention 残差连接LayerNorm MLP 残差连接3.5 分类头取 Encoder 输出序列中 CLS Token 位置的特征向量过一层线性层得到分类 logits。4. 环境准备与最小依赖在动手写代码之前先确认环境。本文的代码基于深度学习框架 PyTorch不依赖额外的第三方库。# 建议使用 Python 3.8 及以上版本 # 安装 PyTorch版本请以官方安装命令为准 pip install torch torchvision验证安装是否成功import torch print(torch.__version__) print(torch.cuda.is_available())需要注意两点本文的代码在 CPU 上也能运行。显存紧张的情况下建议先用小配置跑通逻辑再放大模型规模。如果环境确实没有 PyTorch请先前往 PyTorch 官网根据操作系统、Python 版本和 CUDA 版本生成对应的安装命令不要使用来源不明的安装源。5. Patch Embedding 完整代码实现Patch Embedding 在代码层面有两种实现方式一种是直观的“手动切块 线性映射”另一种是工程上更高效的“Conv2d 一步实现”。两种方式数学上完全等价我建议都掌握。5.1 手动切块 线性映射最直观的实现import torch import torch.nn as nn class PatchEmbeddingNaive(nn.Module): 最直观的 Patch Embedding 实现 1. 把图像切成不重叠的 patch 2. 每个 patch 展平成向量 3. 过线性层映射到 embed_dim def __init__(self, in_channels3, patch_size16, embed_dim768): super().__init__() self.patch_size patch_size # patch 展平后的维度 通道数 × patch 高 × patch 宽 flatten_dim in_channels * patch_size * patch_size self.linear nn.Linear(flatten_dim, embed_dim) def forward(self, x): B, C, H, W x.shape p self.patch_size assert H % p 0 and W % p 0, \ f输入尺寸 {H}x{W} 必须能被 patch_size{p} 整除 # unfold 可以看作“滑动窗口取块”按步长 p 切出 p×p 的块 # 结果 shape: [B, C, H/p, W/p, p, p] x x.unfold(2, p, p).unfold(3, p, p) # 调整维度顺序把空间分块放在前两维 # [B, C, H/p, W/p, p, p] - [B, H/p, W/p, C, p, p] x x.permute(0, 2, 3, 1, 4, 5).contiguous() # 把 C×p×p 展平成一个向量 # [B, H/p, W/p, C*p*p] x x.flatten(3) # 把 H/p 和 W/p 合并成 Patch 数量 N # [B, N, C*p*p] 其中 N H/p × W/p x x.flatten(1, 2) # 线性映射到 embed_dim # [B, N, embed_dim] x self.linear(x) return x手动实现的核心是理解 unfold 的含义。unfold(2, p, p)表示在第 2 维高度方向上以窗口大小为 p、步长为 p 取出块。第 3 维同理。这样一次操作就把整张图切成了 (H/p) × (W/p) 个 Patch。5.2 用 Conv2d 实现工程上最优雅的写法class PatchEmbedding(nn.Module): 工程上最推荐的 Patch Embedding 实现。 用 kernel_sizepatch_size, stridepatch_size 的卷积一步完成 “切块 展平 线性映射”。 def __init__(self, in_channels3, patch_size16, embed_dim768): super().__init__() self.patch_size patch_size # 这一步卷积同时完成了三件事 # 1. kernel_sizepatch_size 相当于取一个 patch 大小的窗口 # 2. stridepatch_size 相当于 patch 之间不重叠 # 3. 输出通道数为 embed_dim相当于每个 patch 线性映射到 embed_dim self.proj nn.Conv2d( in_channels, embed_dim, kernel_sizepatch_size, stridepatch_size ) def forward(self, x): B, C, H, W x.shape assert H % self.patch_size 0 and W % self.patch_size 0, \ f输入尺寸 {H}x{W} 必须能被 patch_size{self.patch_size} 整除 # 卷积输出: [B, embed_dim, H/p, W/p] x self.proj(x) # 展平空间维度: [B, embed_dim, N]N H/p × W/p x x.flatten(2) # 转成序列格式: [B, N, embed_dim] x x.transpose(1, 2) return x为什么 Conv2d 能做到和手动切块等价核心在卷积核的滑动方式。当 kernel_size patch_size、stride patch_size 时卷积核每次正好覆盖一个 Patch且 Patch 之间没有重叠。每个输出位置的值就是该 Patch 内的像素值和卷积核做内积的结果。多个输出通道相当于多个线性映射等价于nn.Linear(flatten_dim, embed_dim)。从性能角度看Conv2d 底层是高度优化的矩阵运算比手动 unfold 之后再过 Linear 更快。因此在实际工程代码如 TIMM、HuggingFace Transformers中基本都是用 Conv2d 实现 Patch Embedding。5.3 两种实现方式的对比对比维度手动切块 LinearConv2d 实现可读性容易理解适合教学简洁高效适合工程计算速度相对较慢更快底层高度优化扩展性改逻辑更直观改 patch_size 只需改参数实际项目较少使用TIMM 等库的标准做法我的建议是入门阶段用手动实现理解原理实际写项目时用 Conv2d 版本。6. 从零手撕 ViT完整 Forward 代码这一节是核心。我会按照“Attention → Block → Encoder → 完整 ViT”的顺序从零搭建 Vision Transformer。6.1 多头自注意力模块class Attention(nn.Module): 多头自注意力模块。 输入输出 shape 相同: [B, N, dim] def __init__(self, dim, num_heads8, qkv_biasFalse, attn_drop0.0): super().__init__() self.num_heads num_heads self.scale (dim // num_heads) ** -0.5 # 一个线性层同时生成 Q、K、V输出维度是 3 * dim self.qkv nn.Linear(dim, dim * 3, biasqkv_bias) self.attn_drop nn.Dropout(attn_drop) self.proj nn.Linear(dim, dim) def forward(self, x): B, N, C x.shape # 生成 QKV并拆成多头 # 先 reshape 成 [B, N, 3, num_heads, C // num_heads] qkv self.qkv(x).reshape(B, N, 3, self.num_heads, C // self.num_heads) # 再 permute 成 [3, B, num_heads, N, C // num_heads] qkv qkv.permute(2, 0, 3, 1, 4) q, k, v qkv[0], qkv[1], qkv[2] # 计算注意力分数: [B, num_heads, N, N] attn (q k.transpose(-2, -1)) * self.scale attn attn.softmax(dim-1) attn self.attn_drop(attn) # 注意力加权求和: [B, num_heads, N, C // num_heads] x attn v # 合并多头: [B, N, C] x x.transpose(1, 2).reshape(B, N, C) x self.proj(x) return x这段代码最需要理解的是维度变换。reshape把 C 维拆成3, num_heads, C // num_heads三部分permute把 QKV 这个维度挪到最前面方便用qkv[0]、qkv[1]、qkv[2]分别取出 Q、K、V。6.2 MLP 与 Transformer Blockclass Mlp(nn.Module): Transformer Block 中的前馈网络。 通常包含两个线性层和一个 GELU 激活。 def __init__(self, in_features, hidden_featuresNone, dropout0.0): super().__init__() hidden_features hidden_features or in_features self.fc1 nn.Linear(in_features, hidden_features) self.act nn.GELU() self.fc2 nn.Linear(hidden_features, in_features) def forward(self, x): x self.fc1(x) x self.act(x) x self.fc2(x) return x class Block(nn.Module): 标准 Transformer Encoder Block: LayerNorm - Attention - 残差连接 - LayerNorm - MLP - 残差连接 def __init__(self, dim, num_heads, mlp_ratio4.0): super().__init__() self.norm1 nn.LayerNorm(dim) self.attn Attention(dim, num_headsnum_heads) self.norm2 nn.LayerNorm(dim) self.mlp Mlp(in_featuresdim, hidden_featuresint(dim * mlp_ratio)) def forward(self, x): # Pre-LN 结构先 LayerNorm 再 Attention残差连接在外层 x x self.attn(self.norm1(x)) x x self.mlp(self.norm2(x)) return x注意这里使用的是 Pre-LN 结构即 LayerNorm 放在 Attention 和 MLP 之前。研究表明Pre-LN 比 Post-LN 更稳定训练时不容易出现梯度爆炸这也是 ViT 论文和 TIMM 库常用做法。Block 里两次残差连接的 shape 完全一致都是[B, N, dim]这是 Transformer 能堆叠几十层的基础。6.3 完整 ViT 模型class VisionTransformer(nn.Module): 完整 Vision Transformer: Patch Embedding - CLS Token - Position Embedding - Encoder - 分类头 def __init__(self, img_size224, patch_size16, in_channels3, num_classes1000, embed_dim768, depth12, num_heads12, mlp_ratio4.0): super().__init__() self.patch_embed PatchEmbedding(in_channels, patch_size, embed_dim) # 计算 Patch 数量例如 224/16 14N 14 × 14 196 num_patches (img_size // patch_size) ** 2 # 可学习的 CLS Tokenshape: [1, 1, embed_dim] self.cls_token nn.Parameter(torch.zeros(1, 1, embed_dim)) # 可学习的位置编码包含 CLS Token 的位置所以需要 1 # shape: [1, num_patches 1, embed_dim] self.pos_embed nn.Parameter(torch.zeros(1, num_patches 1, embed_dim)) self.pos_drop nn.Dropout(0.0) # 堆叠 Transformer Encoder Block self.blocks nn.ModuleList([ Block(embed_dim, num_heads, mlp_ratio) for _ in range(depth) ]) self.norm nn.LayerNorm(embed_dim) # 分类头默认接 ImageNet 1000 类 self.head nn.Linear(embed_dim, num_classes) def forward_features(self, x): B x.shape[0] # 1. Patch Embedding: [B, 3, 224, 224] - [B, 196, 768] x self.patch_embed(x) # 2. 拼接 CLS Token # cls_token: [1, 1, 768] - expand 到 [B, 1, 768] cls_token self.cls_token.expand(B, -1, -1) x torch.cat([cls_token, x], dim1) # [B, 197, 768] # 3. 加入位置编码 x x self.pos_embed # [B, 197, 768] # 4. 过 Transformer Encoder for blk in self.blocks: x blk(x) # shape 始终不变 x self.norm(x) return x def forward(self, x): # 提取特征 x self.forward_features(x) # [B, 197, 768] # 取 CLS Token 的输出特征 x x[:, 0] # [B, 768] # 分类头 x self.head(x) # [B, num_classes] return x6.4 Forward 全流程 shape 变化用一个小型配置跑一次完整前向方便观察每一步 shape# 小型配置方便在 CPU 上跑通流程 model VisionTransformer( img_size224, patch_size16, num_classes10, embed_dim192, depth4, num_heads4, mlp_ratio4.0 ) # 模拟输入2 张 224×224 RGB 图像 x torch.randn(2, 3, 224, 224) # 前向传播 out model(x) print(输出 shape:, out.shape) # torch.Size([2, 10])对应每一阶段的 shape 变化如下阶段Shape说明输入图像[2, 3, 224, 224]B23 通道224×224Patch Embedding 后[2, 196, 192]N196embed_dim192拼接 CLS Token 后[2, 197, 192]序列长度 1加入位置编码后[2, 197, 192]与 pos_embed 相加Encoder 每一层后[2, 197, 192]共 4 层shape 不变取 CLS 特征[2, 192]取第 0 个位置分类输出[2, 10]过 Linear 分类头这个表是整个 ViT 前向过程的骨架。只要你的代码 shape 变化和表一致说明模型结构搭建正确。7. 运行结果与效果验证光看 shape 还不够最好能打印出每一层输出的实际值帮助确认模型是否真的在工作。写一个调试版子类在 forward 过程中插入打印语句class ViTDebug(VisionTransformer): def forward_features(self, x): B x.shape[0] print(f[0] 输入图像: {x.shape}) x self.patch_embed(x) print(f[1] Patch Embedding: {x.shape}) cls_token self.cls_token.expand(B, -1, -1) x torch.cat([cls_token, x], dim1) print(f[2] 拼接 CLS Token: {x.shape}, CLS 值: {x[:, 0, :5]}) x x self.pos_embed print(f[3] 加入位置编码: {x.shape}) for i, blk in enumerate(self.blocks): x blk(x) print(f[4] Block {i1}: {x.shape}) x self.norm(x) print(f[5] 最终 LayerNorm: {x.shape}) return x model_debug ViTDebug( img_size224, patch_size16, num_classes10, embed_dim192, depth4, num_heads4 ) x torch.randn(2, 3, 224, 224) out model_debug(x) print(分类输出:, out.shape)运行后应该看到类似输出[0] 输入图像: torch.Size([2, 3, 224, 224]) [1] Patch Embedding: torch.Size([2, 196, 192]) [2] 拼接 CLS Token: torch.Size([2, 197, 192]) [3] 加入位置编码: torch.Size([2, 197, 192]) [4] Block 1: torch.Size([2, 197, 192]) [4] Block 2: torch.Size([2, 197, 192]) [4] Block 3: torch.Size([2, 197, 192]) [4] Block 4: torch.Size([2, 197, 192]) [5] 最终 LayerNorm: torch.Size([2, 197, 192]) 分类输出: torch.Size([2, 10])判断模型是否搭建成功的标准所有 shape 和上面表格一致。前向传播没有报错。分类输出 logits 不是全零或 NaN。CLS Token 的值在整个过程中持续变化说明它确实在参与注意力计算。如果运行失败先不要急着改模型结构。按下面第 8 节排查。我还可以用一个快速实验验证模型“真的能学”用随机生成的标签做一步反向传播看 loss 是否能下降。这个能确认模型参数确实在更新。import torch.nn.functional as F # 随机生成一对数据和标签仅用于验证梯度流 x torch.randn(4, 3, 224, 224) y torch.randint(0, 10, (4,)) model VisionTransformer( img_size224, patch_size16, num_classes10, embed_dim192, depth4, num_heads4 ) optimizer torch.optim.AdamW(model.parameters(), lr3e-4) loss F.cross_entropy(model(x), y) print(初始 loss:, loss.item()) optimizer.zero_grad() loss.backward() optimizer.step() loss2 F.cross_entropy(model(x), y) print(更新后 loss:, loss2.item())如果 loss 能正常下降说明整个模型从 forward 到 backward 的计算图是完全通的。这一步也是后续接入真实训练流程的起点。8. 常见问题与排查思路手撕 ViT 的过程中有几个问题几乎每个人都会遇到。我把常见问题整理成表方便对照排查。问题现象可能原因排查方式解决方案输入图像尺寸不能整除 patch_sizePatch Embedding 的 assert 拦截查看报错信息中的 H、W、patch_size调整 img_size 或 patch_size或者对输入图像做 resize拼接 CLS Token 后维度报错cls_token 的 embed_dim 与 patch_embed 输出维度不一致打印 patch_embed 输出和 cls_token shape确保 embed_dim 参数在所有模块中保持一致位置编码相加报错pos_embed 的序列长度不是 num_patches1打印 x 和 pos_embed 的 shape检查 num_patches 计算是否正确记得 1注意力分数为 NaN学习率过大或初始化不当打印 attn 的统计值降低学习率检查数据是否需要归一化显存不足模型配置过大、batch_size 过大、序列太长查看 GPU 显存占用减小 embed_dim、depth 或 batch_size训练 loss 不下降没有使用 warmup、学习率不合适、数据没归一化打印梯度范数使用 AdamW warmup学习率设为 3e-4 量级分类准确率远低于 CNN数据量不足、训练轮次不够、缺少数据增强观察训练集和验证集 loss 差距增加训练数据、使用随机裁剪/翻转/MixUp 等增强策略下面挑几个重点问题展开说明。8.1 为什么维度老是差 1很多新手第一次跑通模型后发现 pos_embed 相加时报错原因是“我明明算了 num_patches为什么加不上去”问题在于忘了 CLS Token 也占一个位置。Patch Embedding 输出 N 个 Patch拼接 CLS 后序列长度变成 N1所以位置编码的序列长度也必须是 N1。这个 “1” 是 ViT 里最容易漏掉的地方。8.2 为什么用 Conv2d 当 Patch Embedding 不感觉奇怪吗这是初学者最常问的问题之一。Conv2d 通常是用来提取特征的为什么在这里变成了“切块”工具关键在 kernel_size 和 stride 都等于 patch_size。这会导致卷积核覆盖的区域互不重叠输出的每一个空间位置恰好对应原图的一个 Patch。它做的本质上是一个“不重叠的局部线性映射”和先展平再过线性层完全等价。理解这一点后你会发现 Conv2d 其实是一个精巧的工程实现而不是一个“玄学”操作。8.3 为什么我的 ViT 训练不收敛如果在随机数据上都看不到 loss 明显下降先检查三点数据有没有归一化到 [0,1] 或标准化ViT 对输入尺度比较敏感。学习率是否合理ViT 一般用 AdamW初始学习率在 3e-4 左右并配合 warmup。类别数是否正确如果数据集只有 10 类但分类头输出 1000 类loss 很难下降。8.4 显存不够怎么办ViT 的显存占用明显高于同规模 CNN尤其是自注意力矩阵[B, num_heads, N, N]会占用大量显存。解决办法减小 patch_size 会增大序列长度 N显存会显著上升增大 patch_size 则相反。减小 embed_dim 和 depth 是最直接的降显存手段。减少 batch_size或者使用梯度累积。8.5 为什么 Transformer 在图像上的收敛速度比 CNN 慢ViT 没有 CNN 那种强归纳偏置局部性、平移等变性需要更多数据才能学到同样的规律。在小型数据集上ViT 通常打不过同规模 ResNet。这也是为什么后来出现了 DeiT蒸馏 数据增强、Swin Transformer窗口自注意力 层级结构等改进方案。9. 工程建议与后续学习方向9.1 位置编码不是只有一种选择本文实现的是 ViT 原论文的 1D 可学习位置编码。实际工程中还有几种选择2D 位置编码针对图像的行列坐标分别建模但 ViT 论文实验显示 1D 效果已经够好因为 Transformer 可以通过注意力自己学到空间结构。相对位置编码在注意力分数计算中引入 Patch 之间的相对位置偏移常用于 Swin Transformer 等改进模型。正弦位置编码NLP Transformer 的经典方案不需要训练但图像场景中较少使用。如果要做下游目标检测或分割任务位置编码的选择会产生明显影响。建议先跑通 1D 可学习版本再根据任务需求替换。9.2 优化器与训练策略ViT 训练有两条关键经验优化器推荐 AdamW基础学习率通常设为 3e-4 量级batch size 较大时可适当提高。建议使用 warmup即学习率先从小值线性增长到目标值再按余弦退火衰减。这主要是因为 Transformer 对早期训练不稳定比较敏感。如果显存有限batch size 很小可以考虑使用梯度累积保证等效 batch size 足够大。9.3 数据量与正则化ViT 在小数据集上容易过拟合。如果在 CIFAR-10 这类小数据集上训练建议使用随机裁剪、随机翻转、MixUp、CutMix 等数据增强。增加 Dropout 和 DropPath。预训练后微调这是最有效的方案先在 ImageNet-21k 等大数据集上预训练再在下游任务上做 fine-tune。9.4 从 ViT 走向更复杂的视觉 Transformer跑通本文的 ViT 之后下一步建议按这个顺序深入DeiT用知识蒸馏减少 ViT 对超大数据集的依赖训练策略有参考价值。Swin Transformer引入窗口注意力和层级设计解决 ViT 计算复杂度高、缺乏多尺度特征的问题。MAE用自监督掩码重建方式预训练 ViT展示了 Transformer 在图像上的强大表征能力。CLIPViT 作为视觉编码器和文本编码器对齐统一了图文表示。DINO自监督 ViT得到的特征具有很好的语义分割和物体发现能力。理解 ViT 之后这些模型的很多代码你都能快速看懂因为它们都是“Patch Embedding Transformer Encoder”这个基本骨架的变体。建议你把这篇文章里的完整代码自己敲一遍而不是直接复制。敲的过程中在纸上画出每一步的 shape 变化尤其是 QKV 拆分和注意力矩阵的那几个维度。当你做到不用看代码也能说出[B, num_heads, N, N]的来龙去脉时Transformer 的核心原理就真正属于你了。