企业官网建设流程全解析

看 ViT 源码之前，最难的往往不是 Self-Attention 公式，而是一个更基础的问题：图像是二维的，Transformer 是处理一维序列的，这两者之间到底怎么“翻译”？

很多读者从 NLP 的 Transformer 教程入手，理解词向量、理解 QKV 都很顺利。但一接触到 Vision Transformer（ViT），立刻卡在 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×14=196 个 Patch。序列长度从超过 15 万降到 196，这是质的变化。

这个过程叫 Patch Embedding。它同时完成了三件事：

1. 切块：把图像切成不重叠的 Patch。
2. 展平：把每个 Patch 的像素展平成一维向量。
3. 线性映射：把展平后的向量投影到 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、Value：

• Query（查询）：表示“我在找什么”。
• 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 Token

ViT 参考 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, N+1, embed_dim]的参数矩阵，直接加到输入序列上。注意这里 N+1 是因为多了一个 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_channels=3, patch_size=16, embed_dim=768): 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_size=patch_size, stride=patch_size 的卷积一步完成 “切块 + 展平 + 线性映射”。 """ def __init__(self, in_channels=3, patch_size=16, embed_dim=768): super().__init__() self.patch_size = patch_size # 这一步卷积同时完成了三件事： # 1. kernel_size=patch_size 相当于取一个 patch 大小的窗口 # 2. stride=patch_size 相当于 patch 之间不重叠 # 3. 输出通道数为 embed_dim，相当于每个 patch 线性映射到 embed_dim self.proj = nn.Conv2d( in_channels, embed_dim, kernel_size=patch_size, stride=patch_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 两种实现方式的对比

我的建议是：入门阶段用手动实现理解原理，实际写项目时用 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_heads=8, qkv_bias=False, attn_drop=0.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, bias=qkv_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 Block

class Mlp(nn.Module): """ Transformer Block 中的前馈网络。 通常包含两个线性层和一个 GELU 激活。 """ def __init__(self, in_features, hidden_features=None, dropout=0.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_ratio=4.0): super().__init__() self.norm1 = nn.LayerNorm(dim) self.attn = Attention(dim, num_heads=num_heads) self.norm2 = nn.LayerNorm(dim) self.mlp = Mlp(in_features=dim, hidden_features=int(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_size=224, patch_size=16, in_channels=3, num_classes=1000, embed_dim=768, depth=12, num_heads=12, mlp_ratio=4.0): super().__init__() self.patch_embed = PatchEmbedding(in_channels, patch_size, embed_dim) # 计算 Patch 数量，例如 224/16 = 14，N = 14 × 14 = 196 num_patches = (img_size // patch_size) ** 2 # 可学习的 CLS Token，shape: [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], dim=1) # [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 x

6.4 Forward 全流程 shape 变化

用一个小型配置跑一次完整前向，方便观察每一步 shape：

# 小型配置，方便在 CPU 上跑通流程 model = VisionTransformer( img_size=224, patch_size=16, num_classes=10, embed_dim=192, depth=4, num_heads=4, mlp_ratio=4.0 ) # 模拟输入：2 张 224×224 RGB 图像 x = torch.randn(2, 3, 224, 224) # 前向传播 out = model(x) print("输出 shape:", out.shape) # torch.Size([2, 10])

对应每一阶段的 shape 变化如下：

这个表是整个 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], dim=1) 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 {i+1}: {x.shape}") x = self.norm(x) print(f"[5] 最终 LayerNorm: {x.shape}") return x model_debug = ViTDebug( img_size=224, patch_size=16, num_classes=10, embed_dim=192, depth=4, num_heads=4 ) 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])

判断模型是否搭建成功的标准：

1. 所有 shape 和上面表格一致。
2. 前向传播没有报错。
3. 分类输出 logits 不是全零或 NaN。
4. 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_size=224, patch_size=16, num_classes=10, embed_dim=192, depth=4, num_heads=4 ) optimizer = torch.optim.AdamW(model.parameters(), lr=3e-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 的过程中，有几个问题几乎每个人都会遇到。我把常见问题整理成表，方便对照排查。

下面挑几个重点问题展开说明。

8.1 为什么维度老是差 1

很多新手第一次跑通模型后，发现 pos_embed 相加时报错，原因是：“我明明算了 num_patches，为什么加不上去？”

问题在于忘了 CLS Token 也占一个位置。Patch Embedding 输出 N 个 Patch，拼接 CLS 后序列长度变成 N+1，所以位置编码的序列长度也必须是 N+1。这个 “+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 训练有两条关键经验：

1. 优化器推荐 AdamW，基础学习率通常设为 3e-4 量级，batch size 较大时可适当提高。
2. 建议使用 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 在图像上的强大表征能力。
• CLIP：ViT 作为视觉编码器和文本编码器对齐，统一了图文表示。
• DINO：自监督 + ViT，得到的特征具有很好的语义分割和物体发现能力。

理解 ViT 之后，这些模型的很多代码你都能快速看懂，因为它们都是“Patch Embedding + Transformer Encoder”这个基本骨架的变体。

建议你把这篇文章里的完整代码自己敲一遍，而不是直接复制。敲的过程中，在纸上画出每一步的 shape 变化，尤其是 QKV 拆分和注意力矩阵的那几个维度。当你做到不用看代码也能说出[B, num_heads, N, N]的来龙去脉时，Transformer 的核心原理就真正属于你了。