如果你已经跟着这个系列看到第十七章前面那些 Python 语法、NumPy 操作、PyTorch 张量基础、CNN 图像分类应该都不是陌生词汇了。这一章我安排了一个很多初学者会一直往后拖的内容视觉 Transformer也就是 ViT。我的建议是不要先啃论文直接把它拆开、用 Python 重新实现一个能跑的分类模型比背十遍理论都管用。这一章的目标很具体用一个小型图像数据集从 Patch Embedding 开始手写 ViT训练到能稳定收敛的精度然后整理训练过程中最常见的坑。1. 第十七章为什么没继续堆 CNN而是切到 Transformer1.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. 注意力机制拆解从加权求和到 QKV2.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 上得到的就是当前位置聚合后的新向量。整个过程可以用一张简单的表理解符号叫什么作用Qquery当前要找的信息Kkey每个位置的匹配标签Vvalue被取出的真实特征内容对 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_sizestride 也等于 patch_size这样每个卷积位置正好对应一个非重叠 patch。输出通道数等于 embed_dim每个 patch 就变成了一个向量。代码写起来非常短。import torch import torch.nn as nn class PatchEmbed(nn.Module): def __init__(self, img_size128, patch_size16, in_chans3, embed_dim192): super().__init__() self.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 ) 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_size16那么图片会被切成 8x864 个 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, std0.02) if module.bias is not None: nn.init.zeros_(module.bias) class ViT(nn.Module): def __init__( self, img_size128, patch_size16, num_classes5, embed_dim192, depth6, num_heads4, dropout0.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, std0.02) nn.init.trunc_normal_(self.pos_embed, std0.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), dim1) 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_heads4, qkv_biasFalse, dropout0.0): super().__init__() self.num_heads num_heads self.scale (dim // num_heads) ** -0.5 self.qkv nn.Linear(dim, dim * 3, biasqkv_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_dimNone, dropout0.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_heads4, dropout0.1): super().__init__() self.norm1 nn.LayerNorm(dim) self.attn Attention(dim, num_headsnum_heads, dropoutdropout) self.norm2 nn.LayerNorm(dim) self.mlp Mlp(dim, dim * 4, dropoutdropout) def forward(self, x): x x self.attn(self.norm1(x)) x x self.mlp(self.norm2(x)) return x注意这里用的是 Pre-LN 结构也就是先 LayerNorm再进注意力/前馈网络。相比 Post-LNPre-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(brightness0.2, contrast0.2, saturation0.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 可以先设成 2GPU 显存足够的情况下 pin_memoryTrue 能减少数据传输等待。train_ds datasets.ImageFolder(data/ingredients/train, transformtrain_transform) val_ds datasets.ImageFolder(data/ingredients/val, transformval_transform) train_loader DataLoader(train_ds, batch_size64, shuffleTrue, num_workers2, pin_memoryTrue) val_loader DataLoader(val_ds, batch_size64, shuffleFalse, num_workers2, pin_memoryTrue)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_size128, patch_size16, num_classes5, embed_dim192, depth6, num_heads4, dropout0.1 ).cuda() criterion nn.CrossEntropyLoss() optimizer optim.AdamW(model.parameters(), lr3e-4, weight_decay0.05) total_steps len(train_loader) * 60 scheduler optim.lr_scheduler.OneCycleLR( optimizer, max_lr3e-4, total_stepstotal_steps, pct_start0.1, anneal_strategycos )训练轮数我先给 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(dim1) correct (preds labels).sum().item() total images.size(0) train_acc correct / total print(fepoch{epoch1} loss{total_loss/train_acc:.4f} acc{train_acc:.4f})这里有个容易踩的细节scaler.step 之后如果 loss 是 NaNscaler 会自动跳过这次优化器更新。所以不要自己额外判断 loss 是否为 NaN直接交给 GradScaler 处理就行。验证部分记得加torch.no_grad()和model.eval()否则 dropout 和训练状态不一致验证结果会偏高或偏低。5. 训练日志实测与常见问题的排查顺序5.1 一份看起来正常的训练日志我在一个小型数据集上复现时训练日志大概是这样的epoch12 loss0.68 acc0.76 val_acc0.74 lr5.2e-04 epoch24 loss0.45 acc0.83 val_acc0.80 lr4.1e-04 epoch36 loss0.32 acc0.88 val_acc0.82 lr2.8e-04 epoch48 loss0.24 acc0.92 val_acc0.83 lr1.6e-04 epoch60 loss0.19 acc0.95 val_acc0.84 lr1.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 降到 4embed_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, dim1) preds probs.argmax(dim1)如果你想把结果可视化可以直接在测试集上随机取几张图把预测标签和真实标签打印在一起。不要只盯着准确率数字也要看错误的样本长什么样。很多时候模型预测错了是因为样本本身模糊、遮挡严重或者类别之间视觉上确实很像。推理时整个模型计算量集中在注意力矩阵上。128x128 的输入、patch_size16序列长度只有 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 的实际演进路线。动手跑一遍比看十篇模型解析文章都值得。