如果你已经跟着这个系列看到第十七章，前面那些 Python 语法、NumPy 操作、PyTorch 张量基础、CNN 图像分类，应该都不是陌生词汇了。这一章我安排了一个很多初学者会一直往后拖的内容：视觉 Transformer，也就是 ViT。我的建议是不要先啃论文，直接把它拆开、用 Python 重新实现一个能跑的分类模型，比背十遍理论都管用。这一章的目标很具体：用一个小型图像数据集，从 Patch Embedding 开始手写 ViT，训练到能稳定收敛的精度，然后整理训练过程中最常见的坑。
1. 第十七章为什么没继续堆 CNN，而是切到 Transformer
1.1 CNN 的局部视野与全局建模成本
CNN 之所以在图像任务里流行，是因为卷积核天然具备局部特征提取能力：一个 3x3 卷积只看周围 9 个像素，多个卷积层堆叠后才逐渐扩大感受野。对一个 224x224 的输入，若只用常规卷积，要让最后一层看到全图范围的信息，需要足够深的网络；网络越深，梯度传播、显存占用、调参难度都会跟着上来。
注意力机制的做法完全不同：它一步就把任意两个位置直接关联起来，不需要依赖网络深度去“接力”。在 ViT 中，图像被切成分块后，每一块相当于一个词，任意两块之间都可以直接计算相似度。这也是我从 CNN 转向 Transformer 的核心原因：不是 CNN 不好，而是当数据规模变大、任务变复杂，一个能直接建模全局关系的结构，往往比层层堆叠局部卷积更容易扩展。
对于已经看到第十七章的读者来说，CNN 的卷积、池化、BatchNorm 这些概念应该都已经亲手写过了。此时如果把注意力机制的原理也搞懂，你其实就同时拿到了 NLP、CV、多模态三块领域的公共底座。后面再接触任何带 Transformer 字样的模型，都不会觉得陌生。
1.2 视觉 Transformer 到底改了什么
Transformer 原本是处理序列的，但图像也可以序列化。ViT 的思路很直接：把一张 H x W x C 的图像切成 N 个 patch，每个 patch 展平成向量；再用线性层映射到 embedding 空间；加上位置编码；最后交给标准的 Transformer Encoder。分类时在序列开头加一个 cls_token，用这个 token 对应的输出向量去接全连接分类头。
整个过程不依赖手工设计的卷积感受野，也不依赖图像的局部先验。模型决定关注哪个 patch，完全由注意力权重自己学出来。换句话说，CNN 是先告诉你“相邻像素大概率相关”，Transformer 则是让模型自己从数据里找关系。这个特性让它在大规模数据上更容易扩展，也让结构本身变得非常通用。
ViT 也不是没有代价。它需要足够多的数据或足够强的数据增强，否则很容易欠拟合或者过拟合。小数据集上直接训练 ViT，效果不一定比 CNN 好。这一章里的项目，我会特意把模型设置成 tiny 版本，让它在一个普通单卡环境里也能跑完整个流程，而不是给你一个只能看不能练的大模型。
1.3 本章项目目标与学习路径
这一章的实践目标，是用一个五分类的小型图像数据集，完成一个完整的手写 ViT 分类器。所谓“手写”，不是让你调用某个现成包里的 ViT，而是从 Patch Embedding、位置编码、注意力层、Encoder Block 到分类头，全部用 PyTorch 的 nn.Module 拼起来。
我建议的学习路径分四步：先理解注意力公式；再跟着代码模块把模型搭起来；然后配置数据增强和优化器完成训练；最后对照训练日志排查问题。每一步我都会给出实际可运行的代码片段。这样一套走下来，你不只是“见过 ViT”，而是真正具备了自己修改和调试它的能力。
整个项目代码量不大，核心模型部分不到两百行，适合在 Jupyter 或者普通脚本里逐步执行。接下来我们先从注意力机制讲起，因为如果不理解 Q、K、V，后面看 ViT 代码很容易变成“背代码”，改一个参数都不知道会发生什么。
2. 注意力机制拆解：从加权求和到 QKV
2.1 注意力本质：给特征计算相关度
注意力机制本质上是一个加权求和。给一组特征向量，原本每个特征向量的地位是一样的，但我们希望模型学会给重要位置更高的权重。它既不是魔法，也不是黑盒。
你可以想象一个仓库管理员：仓库里有一堆文件，现在需要找“年终总结”相关资料。注意力机制会为每一份文件计算一个相关度分数，然后按这个分数加权取出内容。相关度越高，权重越大；无关文件虽然也参与了加权，但影响接近于零。ViT 中的每一层 Encoder 其实就在反复做类似的事情：每个 patch 去查询其他所有 patch 的信息，然后把有用信息聚合到自己身上。
这个机制和人类看图的方式也有点像。你看一张图片时，目光不会平均分配给每个像素，而是会重点看边缘、关键物体和文字区域。注意力机制学到的正是这种“重点看哪里”的能力。
2.2 Self-Attention 公式逐项解释
Self-Attention 的公式可以写成：
Attention(Q, K, V) = softmax(QK^T / sqrt(d_k)) V
这里 Q 是 query，表示当前这个位置想要查找什么；K 是 key，表示每个候选位置的特征标识；V 是 value，表示每个候选位置真正携带的内容。点积 QK^T 的作用是计算当前查询和每个候选位置的相似度。
除以 sqrt(d_k) 是为了防止点积结果过大。如果 d_k 很大，点积可能得到一个很大的数值，softmax 进去之后梯度会非常小，模型不容易训练。这个缩放操作在实现时很容易被漏掉，漏掉之后通常表现为训练初期 loss 波动很大。
softmax 负责把相似度分数转成概率分布，让所有权重加起来等于 1。最后再把这些权重作用到 V 上，得到的就是当前位置聚合后的新向量。整个过程可以用一张简单的表理解：
| 符号 | 叫什么 | 作用 |
|---|---|---|
| Q | query | 当前要找的信息 |
| K | key | 每个位置的匹配标签 |
| V | value | 被取出的真实特征内容 |
对 ViT 来说，同一个序列里每个 patch 都会轮流充当 query，去查询其他 patch。所以叫 self-attention，因为查询和被查询的对象来自同一张输入图。
2.3 多头注意力为什么有效
单一注意力只能做一组加权求和，但图像特征之间的关系往往不只一种。比如有的 patch 之间是颜色相近，有的 patch 之间是纹理相似，还有的是空间位置构成的语义关系。如果只用一个注意力头，所有关系都挤在同一组权重里，表达能力会受到限制。
多头注意力的做法很简单：把 embedding 切成 h 份，每一份独立做一次注意力，最后把 h 个结果拼回去。这样每个头可以学到不同侧重点，有的关注局部边缘，有的关注全局轮廓，有的关注物体中心。模型整体表达能力立刻上来了。
在代码实现里，多头注意力通常不会真的写多个独立 Linear 层，而是用一个 qkv 线性层同时生成 Q、K、V，再 reshape 成多头形状。这样计算效率更高，代码也更紧凑。下一节手写 ViT 时，你会看到这个操作的完整写法。
3. 手写一个可运行的 ViT：关键模块与完整类
3.1 Patch Embedding：把图像切成“词”
ViT 的第一步是把图片切成 patch。最直接的做法是手动 reshape 每个 patch，再拉平投影到 embedding 空间。不过 PyTorch 里有个很巧妙的等价实现：用步长等于 kernel_size 的 Conv2d。
卷积核 size 等于 patch_size，stride 也等于 patch_size，这样每个卷积位置正好对应一个非重叠 patch。输出通道数等于 embed_dim，每个 patch 就变成了一个向量。代码写起来非常短。
import torch import torch.nn as nn class PatchEmbed(nn.Module): def __init__(self, img_size=128, patch_size=16, in_chans=3, embed_dim=192): super().__init__() self.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 ) def forward(self, x): x = self.proj(x) # [B, embed_dim, H/p, W/p] x = x.flatten(2) # [B, embed_dim, num_patches] x = x.transpose(1, 2) # [B, num_patches, embed_dim] return x如果输入是 128x128 的图，patch_size=16，那么图片会被切成 8x8=64 个 patch。输出形状应该是 [B, 64, 192]。这一步的输出，和 NLP 里“一句话切成若干个词向量”是同一层意思。
注意：patch_size 越小，patch 数量越多，序列越长，注意力矩阵越大，显存和训练时间会明显上涨。小 patch 确实可能带来细粒度优势，但代价很高。先用 16 起步是性价比最高的选择。
3.2 cls_token 与位置编码
图片变成 patch 序列后，还缺两样东西：分类令牌和位置编码。
分类令牌的做法是在序列开头额外放一个可学习的向量，叫 cls_token。这个 token 不承载任何具体 patch 的信息，但它会在所有 Encoder 层中和所有 patch 做注意力交互，所以训练完成后，它对应的输出可以看作整个图像的全局表示。分类头只需要接在这个向量上。
位置编码则解决顺序问题。patch 被切出来之后，模型默认是不认识它们之间顺序的。与 CNN 不同，Transformer 没有内置的邻接关系概念，必须显式告诉它每个 patch 在什么位置。ViT 里常用的是可学习位置编码，也就是一个形状为 [1, num_patches + 1, embed_dim] 的参数，和 token 序列加在一起。
def init_vit_weights(module): if isinstance(module, nn.Linear): nn.init.trunc_normal_(module.weight, std=0.02) if module.bias is not None: nn.init.zeros_(module.bias) class ViT(nn.Module): def __init__( self, img_size=128, patch_size=16, num_classes=5, embed_dim=192, depth=6, num_heads=4, dropout=0.1 ): super().__init__() self.patch_embed = PatchEmbed(img_size, patch_size, 3, embed_dim) num_patches = self.patch_embed.num_patches self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim)) self.pos_embed = nn.Parameter(torch.zeros(1, num_patches + 1, embed_dim)) nn.init.trunc_normal_(self.cls_token, std=0.02) nn.init.trunc_normal_(self.pos_embed, std=0.02) self.pos_drop = nn.Dropout(dropout) self.blocks = nn.Sequential(*[ EncoderBlock(embed_dim, num_heads, dropout) for _ in range(depth) ]) self.norm = nn.LayerNorm(embed_dim) self.head = nn.Linear(embed_dim, num_classes) self.apply(init_vit_weights) def forward(self, x): B = x.shape[0] x = self.patch_embed(x) cls_tokens = self.cls_token.expand(B, -1, -1) x = torch.cat((cls_tokens, x), dim=1) x = x + self.pos_embed x = self.pos_drop(x) x = self.blocks(x) x = self.norm(x) return self.head(x[:, 0])这里 num_patches + 1，是因为序列前面多了一个 cls_token。位置编码的数量必须和序列长度一致，否则x + self.pos_embed这一步直接就会报错。这是手写 ViT 时最容易出现的维度问题。
3.3 注意力层、Encoder Block 与分类头
接下来是关键的头重脚轻部分：Attention 类。实现时会把 Q、K、V 合并成一个线性层，然后 reshape 成多头格式。为了减少代码绕来绕去，我直接用 einsum 还是普通 reshape 都可以，下面这份代码用直观的 reshape，方便你一步步看形状变化。
class Attention(nn.Module): def __init__(self, dim, num_heads=4, qkv_bias=False, dropout=0.0): super().__init__() self.num_heads = num_heads self.scale = (dim // num_heads) ** -0.5 self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias) self.proj = nn.Linear(dim, dim) self.drop = nn.Dropout(dropout) def forward(self, x): B, N, C = x.shape qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, C // self.num_heads) qkv = qkv.permute(2, 0, 3, 1, 4) q, k, v = qkv[0], qkv[1], qkv[2] attn = (q @ k.transpose(-2, -1)) * self.scale attn = attn.softmax(dim=-1) attn = self.drop(attn) x = (attn @ v).transpose(1, 2).reshape(B, N, C) x = self.proj(x) return x这里 q 和 k 的最后一维是 head_dim，也就是 dim 除以 num_heads。scale 用的是 head_dim 的 -0.5 次方，而不是 dim 直接。很多初学实现容易把 scale 算错，导致训练不稳定。
Encoder Block 需要两个子层：注意力层和前馈网络层。前馈网络就是两个 Linear 中间夹一个 GELU 激活函数，扩展比例通常是 embedding 维度的 4 倍。每个子层外面再接残差连接和 LayerNorm，结构如下。
class Mlp(nn.Module): def __init__(self, dim, hidden_dim=None, dropout=0.0): super().__init__() hidden_dim = hidden_dim if hidden_dim is not None else dim * 4 self.fc1 = nn.Linear(dim, hidden_dim) self.act = nn.GELU() self.fc2 = nn.Linear(hidden_dim, dim) self.drop = nn.Dropout(dropout) def forward(self, x): x = self.fc1(x) x = self.act(x) x = self.fc2(x) x = self.drop(x) return x class EncoderBlock(nn.Module): def __init__(self, dim, num_heads=4, dropout=0.1): super().__init__() self.norm1 = nn.LayerNorm(dim) self.attn = Attention(dim, num_heads=num_heads, dropout=dropout) self.norm2 = nn.LayerNorm(dim) self.mlp = Mlp(dim, dim * 4, dropout=dropout) def forward(self, x): x = x + self.attn(self.norm1(x)) x = x + self.mlp(self.norm2(x)) return x注意这里用的是 Pre-LN 结构，也就是先 LayerNorm，再进注意力/前馈网络。相比 Post-LN，Pre-LN 在训练稳定性上更好，尤其适合从头训练一个小模型。残差连接的作用是让梯度能够直接回传，让每个 Block 只是在原来的特征上做增量更新。
到这一步，核心模型已经全部手写完毕。总共只有三个核心类：PatchEmbed、Attention、EncoderBlock，加上一个 ViT 主类。跑一下模型前向，输出形状就是 [batch_size, num_classes]。
4. 数据准备与训练配置：稳定收敛比模型创新更重要
4.1 数据增强与数据集加载
ViT 没有 CNN 那么强的归纳偏置，所以对数据增强更敏感。我用的数据集是一个小型食材分类数据集，图片尺寸统一处理成 128x128，类别数 5。虽然数据量不大，但只要增强策略到位，还是能训练到一个不错的效果。
训练集增强我通常这样设置：随机裁剪、随机水平翻转、颜色抖动。验证集不使用随机增强，只做 Resize 和 CenterCrop，保证评估结果稳定。代码和普通 torchvision 流程没有区别。
from torch.utils.data import DataLoader from torchvision import datasets, transforms train_transform = transforms.Compose([ transforms.RandomResizedCrop(128, scale=(0.7, 1.0)), transforms.RandomHorizontalFlip(), transforms.ColorJitter(brightness=0.2, contrast=0.2, saturation=0.2), transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) ]) val_transform = transforms.Compose([ transforms.Resize(128), transforms.CenterCrop(128), transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) ])如果你的数据集目录结构是每个类别一个文件夹，可以直接用 ImageFolder 加载。加载时 num_workers 可以先设成 2，GPU 显存足够的情况下 pin_memory=True 能减少数据传输等待。
train_ds = datasets.ImageFolder("data/ingredients/train", transform=train_transform) val_ds = datasets.ImageFolder("data/ingredients/val", transform=val_transform) train_loader = DataLoader(train_ds, batch_size=64, shuffle=True, num_workers=2, pin_memory=True) val_loader = DataLoader(val_ds, batch_size=64, shuffle=False, num_workers=2, pin_memory=True)4.2 优化器、学习率与训练轮数
ViT 的默认优化器组合，在 PyTorch 生态里基本就是 AdamW。官方 ViT 论文和后续很多复现都用 AdamW，配合 weight_decay 做正则化。学习率一般从 1e-4 到 3e-4 之间开始试，太大容易训练震荡，太小则收敛很慢。
因为从头训练 Transformer 往往需要 warmup，我直接用 OneCycleLR 调度器，它会在前 10% 的迭代里把学习率逐步升到 max_lr，再用余弦退火降下来。这样做的好处是不用自己写 warmup 和 cosine 两段逻辑，一个调度器全部搞定。
import torch.optim as optim model = ViT( img_size=128, patch_size=16, num_classes=5, embed_dim=192, depth=6, num_heads=4, dropout=0.1 ).cuda() criterion = nn.CrossEntropyLoss() optimizer = optim.AdamW(model.parameters(), lr=3e-4, weight_decay=0.05) total_steps = len(train_loader) * 60 scheduler = optim.lr_scheduler.OneCycleLR( optimizer, max_lr=3e-4, total_steps=total_steps, pct_start=0.1, anneal_strategy="cos" )训练轮数我先给 60。epoch 不是越多越好，很多情况下模型在 40 轮以后就开始出现过拟合趋势。与其盲目加轮数，不如把模型做得小一点、数据增强做好一点。
4.3 训练主循环与梯度缩放
如果 GPU 支持混合精度，建议直接用自动混合精度。Transformer 对显存的消耗比普通 CNN 大，AMP 能把显存占用明显降下来，训练速度也会更快。PyTorch 提供的做法是 GradScaler 配合 autocast。
scaler = torch.cuda.amp.GradScaler() for epoch in range(60): model.train() total_loss = 0.0 correct = 0 total = 0 for images, labels in train_loader: images = images.cuda() labels = labels.cuda() optimizer.zero_grad() with torch.cuda.amp.autocast(): logits = model(images) loss = criterion(logits, labels) scaler.scale(loss).backward() scaler.step(optimizer) scaler.update() scheduler.step() total_loss += loss.item() * images.size(0) preds = logits.argmax(dim=1) correct += (preds == labels).sum().item() total += images.size(0) train_acc = correct / total print(f"epoch={epoch+1} loss={total_loss/train_acc:.4f} acc={train_acc:.4f}")这里有个容易踩的细节：scaler.step 之后，如果 loss 是 NaN，scaler 会自动跳过这次优化器更新。所以不要自己额外判断 loss 是否为 NaN，直接交给 GradScaler 处理就行。
验证部分记得加torch.no_grad()和model.eval()，否则 dropout 和训练状态不一致，验证结果会偏高或偏低。
5. 训练日志实测与常见问题的排查顺序
5.1 一份看起来正常的训练日志
我在一个小型数据集上复现时，训练日志大概是这样的：
epoch=12 loss=0.68 acc=0.76 val_acc=0.74 lr=5.2e-04 epoch=24 loss=0.45 acc=0.83 val_acc=0.80 lr=4.1e-04 epoch=36 loss=0.32 acc=0.88 val_acc=0.82 lr=2.8e-04 epoch=48 loss=0.24 acc=0.92 val_acc=0.83 lr=1.6e-04 epoch=60 loss=0.19 acc=0.95 val_acc=0.84 lr=1.0e-04这份日志是健康的：训练 loss 持续下降，验证准确率也同步上升，训练和验证的差距到最后约 0.1 左右。对一个五分类小型数据集，ViT 能达到 0.84 验证准确率已经说明结构本身没写错，注意力机制确实在工作。
如果验证准确率从头到尾都不动，比如训练 loss 也在 1.6 附近徘徊，那大概率不是数据问题，而是模型前向或者 loss 计算出了问题。这种情况下不要盲目调学习率，先跑一个 batch，打印 logits 和 loss，检查输出形状和数值范围。
5.2 训练集很好、验证集很差怎么办
一个典型现象是训练集准确率已经到 0.95，验证集只有 0.55。这说明模型记住了训练集细节，没有学到泛化特征。ViT 因为结构灵活，拟合能力强，在小数据上非常容易过拟合。
排查顺序应该是先看数据增强是否足够。如果训练集加载后没有随机裁剪、翻转和颜色扰动，模型很快就会背住训练集图片。再往下看正则化，weight_decay 可以从 0.05 提高到 0.1。也可以增加 dropout。还有一个简单做法是减少模型尺寸，把小模型的 depth 从 6 降到 4，embed_dim 从 192 降到 128。
我个人的经验是：如果模型过拟合，优先调整数据增强，而不是立刻改模型。增强的收益往往比调正则化参数更明显。
5.3 三个我优先调整的设置
先列一个速查表格，方便你遇到问题时直接对照：
| 现象 | 排查优先级 | 可能原因 | 处理方式 |
|---|---|---|---|
| 显存溢出 | 高 | batch_size 太大、patch_size 太小 | 降低 batch，开启 AMP，或使用梯度累积 |
| loss 不下降 | 高 | 学习率不合适、数据归一化错误 | 先用一个 batch 过拟合，再调学习率 |
| 训练好验证差 | 中 | 数据增强不足、模型过拟合 | 增强增强策略，加大 weight_decay |
| 训练速度很慢 | 中 | num_workers 太低、patch 数太多 | 提高 workers，加大 patch_size |
在我实际训练经验里，最容易被忽略的是“用一个 batch 过拟合”这个 sanity check。如果模型连一个 batch 都拟合不了，那问题一定在模型或数据处理，而不是超参数。只要模型能在一个 batch 上把 loss 降到接近 0，再恢复完整训练，成功率会高很多。
6. 测试推理与后续扩展：从 tiny ViT 到可落地模型
6.1 推理阶段的代码与细节
训练完成后，推理阶段要记得切换模型状态。model.eval()会关闭 dropout 和 batch norm 的训练行为，torch.inference_mode()比torch.no_grad()在某些场景下还能再快一点。推理代码如下。
model.eval() with torch.inference_mode(): images = images.cuda() logits = model(images) probs = torch.softmax(logits, dim=1) preds = probs.argmax(dim=1)如果你想把结果可视化，可以直接在测试集上随机取几张图，把预测标签和真实标签打印在一起。不要只盯着准确率数字，也要看错误的样本长什么样。很多时候模型预测错了，是因为样本本身模糊、遮挡严重，或者类别之间视觉上确实很像。
推理时整个模型计算量集中在注意力矩阵上。128x128 的输入、patch_size=16，序列长度只有 65，所以单张推理很快。如果你要部署到 CPU 或移动端，可以考虑把模型宽度继续缩小，或者导出成通用推理格式再优化。
6.2 从 ViT 延伸出去的三个方向
第一个方向是目标检测和分割。ViT 可以直接作为骨干网络，也可以把特征图变成序列输入到检测头里。理解了 cls_token 和 patch 序列的关系后，再看检测模型，你会发现它只是把“全局分类”改成了“局部定位”。
第二个方向是多模态。用同一个 Transformer 结构，一边编码图像 patch 序列，一边编码文本 token 序列，让图像和文本的特征空间对齐。图像编码器部分基本就是这一章写的 ViT，只是最后不一定用分类头，而是输出一个特征向量。所以这一章搭建的代码，其实已经是很多多模态模型的视觉侧底子。
第三个方向是轻量化。原始 ViT 参数不小，但可以改 patch、减少 depth、缩小 embed_dim，得到 tiny 版本。也可以把注意力换成线性注意力或者局部注意力，减少长序列计算。先跑通这一章的 tiny ViT，再做这些替换会容易很多。
6.3 部署准备与个人实操体会
部署时最基础的三件事是：固定随机种子、固定输入尺寸、导出模型前先做一次验证集精度确认。不要训练完直接导出，最好先用测试集测一遍，确认model.eval()状态下的精度和训练时验证精度一致，再进入部署流程。
如果部署环境是 GPU 服务，可以先把模型转成半精度，既能省显存也能提升吞吐。如果部署环境是 CPU，需要重点确认输入图片预处理是否和训练一致。Resize 尺寸、Normalize 用的 mean/std，任何一个地方不一致，部署效果都会明显变差。
最后说点个人体会。我每次拿到新数据集，会先用一个很小的 ViT 快速试跑，而不是一上来就堆大模型。因为小模型能让问题暴露得很快——如果小模型都不收敛，大模型只会更难调。你把这套 tiny ViT 跑通之后，再往 depth、embed_dim、patch_size 三个方向放大，就是很多视觉 Transformer 的实际演进路线。动手跑一遍，比看十篇模型解析文章都值得。
