Vision TransformerViT这几年在图像分类领域已经不是新鲜词了但很多人第一次接触它的时候都会被“把图片切成 patch 再当序列处理”这个思路绕晕。老实说我第一次看 ViT 论文的时候也琢磨了很久——明明 CNN 在处理图像上已经那么成熟了为什么还要用一个原本给 NLP 用的 Transformer 来做图像分类这篇文章我就从实际干活的角度把 ViT 的原理、实现、训练避坑一次性讲清楚。内容覆盖从 patch embedding 到 encoder 的完整流程也包含我自己的 PyTorch 实现经验和调参心得适合已经会用 CNN 做分类、想搞懂 ViT 内部机制的人群也适合准备在项目里实际用 ViT 的开发者。1. ViT 到底解决了什么问题1.1 从 CNN 到 Transformer 的思路跃迁CNN 处理图像的核心假设是局部性locality和空间不变性translation invariance。卷积核只在局部区域滑动通过堆叠层数逐步扩大感受野。这个设计在 ImageNet 这种大规模数据集上被验证了十几年非常可靠。但它的缺点是明显的要捕获图像中相距很远的两个区域的关系必须把网络做得非常深或者依赖空洞卷积、特征金字塔这些额外设计。长距离依赖建模对 CNN 来说始终是个绕弯子的事情。Transformer 的思路则完全不同。自注意力机制让任意两个位置之间可以直接建立联系一步到位不存在感受野慢慢扩大的过程。ViT 把图像切成一个个 patch把每个 patch 当作一个“词”然后用标准 Transformer encoder 去建模 patch 与 patch 之间的全局关系。换句话说图像被“翻译”成了一段序列剩下的交给注意力机制去发现哪里重要、哪里和哪里相关。1.2 一句话概括 ViT 的做法ViT 的工作流程可以压缩成四步先把图像切成固定大小的 patch比如 16x16然后通过线性投影把每个 patch 变成一个向量接着加上位置编码让模型知道每个 patch 在图像里的位置然后送入多层 Transformer encoder让 patch 之间充分交换信息最后取出序列最前面的 class token接一个分类头输出类别概率。整个模型几乎原封不动地搬了 NLP 里 BERT 的架构唯一的图像相关改动就是最前面的 patch embedding。这也是 ViT 最“反直觉”的地方一个完全不懂图像先验结构的模型在足够大的数据集上训练之后居然能媲美甚至超过精心设计的 CNN。这也是为什么很多人说 ViT 证明了“先验知识不是必须的数据量才是”。1.3 适合谁来读这篇如果你已经用 ResNet、EfficientNet 做过图像分类想了解为什么 ViT 在大规模数据上表现更好或者你手头有个分类任务想在项目里尝试 ViT 但不知道从哪下手又或者你在读 ViT 论文时对 patch embedding、position embedding、class token 这些概念似懂非懂想看到一个能跑通的实现——这篇文章就是写给你的。我会尽量不堆公式用代码和实验结果说话。2. ViT 架构逐层拆解2.1 Patch Embedding把图片切成“单词”Patch embedding 是整个 ViT 里唯一跟图像强相关的部分也是理解 ViT 的关键一步。假设输入图像是 224x224x3patch size 设为 16那么图像会被切成(224/16) × (224/16) 14 × 14 196个 patch。每个 patch 原本是一个 16x16x3 768 维的张量经过一层线性投影后变成一个长度为 D 的向量。这个线性投影在实现上就是一个卷积层kernel size 和 stride 都等于 patch size。这一步很多人第一次看代码时会困惑明明论文说的是“线性投影”为什么代码里是 Conv2d因为Conv2d(in_channels3, out_channelsD, kernel_size16, stride16)的数学本质就是把每个 16x16x3 的窗口展平后做矩阵乘法和 Linear 等价但用卷积实现效率高得多还顺便处理了 batched input。2.2 Position Embedding给 patch 加上顺序感自注意力机制本身是“排列不变”的也就是说把 196 个 patch 的顺序打乱模型输出的结果完全一样。这对图像分类来说是不可接受的所以必须把位置信息显式注入。ViT 的默认做法是给每个 patch 向量直接加上一个可学习的 position embedding。这个 embedding 的维度是(num_patches 1, D)多出来的 1 是给 class token 用的。它有几种备选方案用 sin-cos 的固定编码、用 2D 的相对位置编码或者干脆用可学习的 1D 编码。ViT 论文里做了对比实验结论是可学习 1D 编码在 ImageNet 上效果最好而且实现最简单。我自己的实验也验证了这一点——在中小规模数据集上几种位置编码的差距很小没必要在这个地方花太多心思。2.3 Transformer Encoder 内部结构ViT 用的 encoder block 和原始 Transformer 论文里的 encoder 几乎一致每个 block 由两部分组成多头自注意力MSA和多层感知机MLP中间用 LayerNorm 和残差连接串起来。有个细节值得注意ViT 用的是 Pre-LN 结构也就是先做 LayerNorm 再做 Attention这和原始 Transformer 的 Post-LN 不一样。Pre-LN 的梯度传播更稳定训练大模型时不容易爆炸。MLP 部分有一个固定的通道扩展比例通常是 4 倍。比如 embedding 维度是 768MLP 内部的隐藏层就是 3072。激活函数用的是 GELU比 ReLU 更平滑在 Transformer 类模型里是标配。self-attention 的计算过程说起来不复杂每个 patch 向量经过三个不同的线性投影得到 Query、Key、Value然后用 Query 和所有 Key 做点积除以维度的平方根做缩放再过 softmax 得到注意力权重最后用权重对 Value 加权求和。多头就是把这个过程分成 H 个头并行做每个头关注不同的关系模式最后拼接起来再投影一次。2.4 Classification Head 与 class token 的设计分类任务需要一个固定长度的输出向量。ViT 的做法是在 patch 序列的最前面拼接一个可学习的 class token它也是一个 D 维向量初始值随机。这个 token 和其他 patch 一样经过所有 encoder block最终在输出层取出它的向量接一个 LayerNorm 和一个全连接层输出类别数。为什么要单独搞一个 class token而不是直接对所有 patch 的输出去全局池化论文的解释是class token 相当于模型自己学习出来的一个“内容聚合器”它在训练中逐渐学会从其他 patch 那里收集全局信息。从实验结果看class token 和全局平均池化GAP的效果差不多但 class token 和预训练模型的兼容性更好因为预训练权重就是按这个结构训练出来的。如果你要用开源预训练权重做微调建议保留 class token 结构别擅自改成 GAP。3. 从零实现 ViT 关键模块3.1 Patch Embedding 的代码实现这一段我直接放可以跑的 PyTorch 代码逻辑和官方实现一致但去掉了不必要的抽象方便你读懂每一行。import torch import torch.nn as nn class PatchEmbed(nn.Module): 图像转 patch 序列用卷积实现线性投影 def __init__(self, img_size224, patch_size16, in_chans3, embed_dim768): super().__init__() self.img_size img_size 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: (B, 3, 224, 224) - (B, 768, 14, 14) x self.proj(x) # 展平成序列flatten 后调换维度 # (B, 768, 14, 14) - (B, 196, 768) x x.flatten(2).transpose(1, 2) return x注意flatten(2)把最后两个维度展平得到(B, 768, 196)再用transpose(1, 2)换成(B, 196, 768)这样每一行就是一个 patch 的 embedding 向量正好符合 Transformer 输入(batch, seq_len, hidden_dim)的约定。3.2 Position Embedding 与 class token 的拼装在把 patch 序列送进 encoder 之前要先拼接 class token、加上位置编码。实现代码如下class VisionTransformer(nn.Module): def __init__(self, img_size224, patch_size16, in_chans3, num_classes1000, embed_dim768, depth12, num_heads12, mlp_ratio4.0): super().__init__() self.patch_embed PatchEmbed(img_size, patch_size, in_chans, 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.pos_embed, std0.02) nn.init.trunc_normal_(self.cls_token, std0.02) # 完整 encoder 由多个 Block 组成见 3.3 节 self.blocks nn.ModuleList([ Block(embed_dim, num_heads, mlp_ratio) for _ in range(depth) ]) self.norm nn.LayerNorm(embed_dim) self.head nn.Linear(embed_dim, num_classes) def forward(self, x): B x.shape[0] x self.patch_embed(x) # (B, 196, 768) cls_tokens self.cls_token.expand(B, -1, -1) x torch.cat([cls_tokens, x], dim1) # (B, 197, 768) x x self.pos_embed for blk in self.blocks: x blk(x) x self.norm(x) cls_out x[:, 0] # 取出 class token 对应的输出 return self.head(cls_out)初始化这里有个小细节cls_token和pos_embed用trunc_normal_初始化标准差设为 0.02这是 Transformer 类模型通用的初始化策略。千万别用nn.init.zeros_去初始化 pos_embed否则前向传播时所有 patch 共享相同的位置信号模型一开始就分不清位置收敛会明显变慢。3.3 Encoder Block 实现Encoder block 是 ViT 的核心计算单元实现了自注意力、MLP、残差连接和 LayerNorm。代码比 patch embedding 稍长但结构非常清晰class Attention(nn.Module): def __init__(self, dim, num_heads12): super().__init__() self.num_heads num_heads self.head_dim dim // num_heads self.scale self.head_dim ** -0.5 self.qkv nn.Linear(dim, dim * 3) self.proj nn.Linear(dim, dim) def forward(self, x): B, N, C x.shape qkv self.qkv(x).reshape(B, N, 3, self.num_heads, self.head_dim) qkv qkv.permute(2, 0, 3, 1, 4) q, k, v qkv.unbind(0) # 每个都是 (B, num_heads, N, head_dim) attn (q k.transpose(-2, -1)) * self.scale attn attn.softmax(dim-1) x (attn v).transpose(1, 2).reshape(B, N, C) return self.proj(x) class MLP(nn.Module): def __init__(self, dim, mlp_ratio4.0): super().__init__() hidden_dim int(dim * mlp_ratio) self.fc1 nn.Linear(dim, hidden_dim) self.act nn.GELU() self.fc2 nn.Linear(hidden_dim, dim) def forward(self, x): return self.fc2(self.act(self.fc1(x))) class Block(nn.Module): def __init__(self, dim, num_heads, mlp_ratio4.0): super().__init__() self.norm1 nn.LayerNorm(dim) self.attn Attention(dim, num_heads) self.norm2 nn.LayerNorm(dim) self.mlp MLP(dim, mlp_ratio) def forward(self, x): x x self.attn(self.norm1(x)) x x self.mlp(self.norm2(x)) return x这里的缩放因子self.scale self.head_dim ** -0.5就是论文里的1/sqrt(d_k)。它在数值上非常重要如果不做缩放当 head_dim 比较大时Query 和 Key 的点积数值会很大softmax 之后梯度趋近于零模型很难训练。我在实验里刻意去掉过一次这个缩放结果是一模一样的代码loss 死活降不下去加回来之后马上正常。3.4 模型规模选择与参数量估算ViT 有几种常用的规格我整理了一个速查表方便你在实际项目中选型配置embed_dimdepthnum_headsMLP隐藏层参数量约ViT-Tiny192123768550万ViT-Small38412615362200万ViT-Base768121230728600万ViT-Large1024241640963.07亿ViT-Base 是论文里最常用的配置对应vit_base_patch16_224。如果你只是做实验验证想法建议从 ViT-Tiny 或 Small 入手在 CPU 上也能勉强跑通 forward显存压力小很多。我最初直接上 ViT-Base 在单张消费级显卡上跑batch size 稍微调大一点就显存溢出换成 Tiny 之后整个调试周期缩短了不止一半。4. 训练配置与实验经验4.1 数据规模是 ViT 的命门ViT 和 CNN 最大的不同在于对数据量的敏感度。ViT 论文的核心结论是在 ImageNet-1k128万张图上直接从头训练ViT 打不过同级别的 ResNet但在 ImageNet-21k1400万张图或 JFT-300M3亿张图上预训练之后ViT 的表现反超 CNN。原因是 ViT 几乎没有图像先验所有归纳偏置都得靠数据自己学出来。这给实际项目带来的启示很直接如果你的数据集只有几千张图别指望从头训练 ViT 能出好效果。正确的做法是用在大规模数据集上预训练好的权重做微调。比如你可以在 HuggingFace 或官方仓库里加载 ImageNet-21k 预训练的 ViT 权重然后在自己的数据上微调通常几百到几千张图就能获得不错的结果。如果你的条件和数据都很紧张那就老实回归 ResNet 系模型没必要为了“跟上潮流”硬上 ViT。4.2 Optimizer、学习率与 warmup 策略ViT 训练时的优化器选择和 CNN 时代有一个显著区别主流做法是用 AdamW而不是带动量的 SGD。我自己对比过在相同数据上 AdamW 的收敛速度明显更快最终精度也略高。学习率方面ViT 对学习率非常敏感论文里用的基础学习率是 0.003配合 linear warmup cosine decay 的调度。Warmup 这一步千万不能省。ViT 的训练初期非常脆弱直接上大学习率很容易让 loss 冲到 NaN。我的习惯是 warmup 步数设为总步数的 5% 到 10%比如总共训练 100 个 epoch前 5 到 10 个 epoch 做线性 warmup。另外weight decay 建议设大一点ViT 的默认值是 0.05 到 0.1远高于 ResNet 常用的 1e-4。这个参数对最终精度的影响很明显我试过把 0.05 改成 0.01同样训练 100 epoch验证集准确率掉了将近 1 个点。4.3 数据增强与正则化技巧ViT 训练还有一个和 CNN 明显不同的地方它对数据增强和正则化的依赖性更强。ViT 论文里用了一整套增强组合包括 RandAugment、Mixup、CutMix、随机擦除等。如果你只是简单做随机裁剪加水平翻转在中小数据集上很容易过拟合。我的经验是ViT 微调场景下RandAugment Mixup 这个组合性价比最高。如果你想控制实验变量可以先不加 CutMix因为 CutMix 和 Mixup 同时使用时需要仔细调参数否则容易互相干扰。DropPathStochastic Depth也是一个值得加的正则化项它随机丢弃某些 block 的残差分支对 ViT 这类深层模型很有效果通常设 0.1 左右。还有一个容易被忽略的细节是 Label Smoothing默认设 0.1能有效缓解模型过度自信的问题提升泛化能力。5. 常见问题与排查实录5.1 训练不收敛或者 loss 直接变 NaN这个是我被问得最多的问题。先检查几个最常见的坑第一学习率是否过大ViT 建议从 1e-4 这个量级起步如果用到 1e-3 还得确保 warmup 足够第二检查数据归一化ViT 需要和预训练时一致的 mean/std如果用 ImageNet 预训练权重输入必须按 ImageNet 的归一化参数来预处理否则 patch 数值分布完全对不上第三检查模型输出层之前的 LayerNorm 是否存在这个self.norm在论文结构里是明确有的删掉之后精度会掉严重时甚至不收敛。还有一个我踩过的坑是混合精度训练时的 loss 溢出。如果你用 AMP建议在 loss 缩放上做文章或者干脆先关掉 AMP 跑一个完整 epoch 确认模型本身没问题再开启混合精度。我遇到过精度初始化不一致导致复现不了实验的情况排查到最后才发现是 AMP 的 loss scale 策略不同。5.2 显存爆炸怎么处理ViT 的显存消耗比同参数量级的 CNN 大。原因是 self-attention 需要计算(B, H, N, N)的注意力矩阵N 是 patch 数加 1。以 ViT-Base 为例224x224 输入对应 197 个 token196x196 的注意力矩阵不算大但如果你把输入图调到 384x384N 就变成 577注意力矩阵大小翻了将近 9 倍。这解释了为什么 ViT 在更高分辨率上推理时会突然显存吃紧。几个实用的缓解手段一是减小 batch size配合梯度累积gradient accumulation保持等效 batch size 不变二是用torch.utils.checkpoint做激活重计算虽然会拖慢训练速度但显存能省一半以上三是如果必须做高分辨率推理可以考虑用窗口注意力或者用 Swin Transformer 这类层级化设计它们在处理高分辨率输入时更有优势。5.3 在中小数据集上过拟合严重如果你在小数据集上直接从零训练 ViT很快就会发现训练集 loss 降得很低但验证集准确率上不去这就是过拟合。解决路径有两条优先级从高到低排序第一条是用预训练权重做微调这是最有效的方案没有之一第二条是加大正则化力度包括增强 DropPath 比例、提高 weight decay、加入更强的数据增强。我也试过一些物理层面的技巧比如冻结前几层 encoder block 的权重只微调后面几层和分类头。这种方式在数据量极少比如每类只有几十张图的时候特别有用相当于把 ViT 当特征提取器用只训练最上层。实测下来在大概 2000 张图的二分类任务上这种“部分微调”比“全量微调”稳定得多不容易把预训练学到的通用特征破坏掉。5.4 和 CNN 对比时怎么做到公平不少人在实验里会遇到一个尴尬的情况ViT 和 ResNet 对比时表现并没有优势甚至更差。这时候先别急着下结论检查一下你的对比设置是否公平。第一训练 epoch 数是否足够ViT 通常需要 200 到 300 个 epoch 才能发挥出完整能力用 60 个 epoch 对比对 ViT 不公平第二数据增强和正则化是否按各自最优配置来设置ViT 对增强的依赖更强第三是否使用了预训练权重这个差异比架构差异大得多。如果是推理速度对比要特别注意 ViT 的高分辨率扩展性问题同时也要考虑实际部署平台。在 GPU 上 ViT 的矩阵运算非常高效但在移动端或嵌入式设备上纯卷积的 CNN 往往更占优势。选型的时候别只盯着精度数字还是要把推理延迟、显存占用、模型大小这些工程指标一起放进表格里综合评估。6. 两个亲测有效的使用建议最后分享两个我实际用下来很有价值的小技巧。第一个是关于微调时的学习率分层分类头用相对较大的学习率比如 1e-3而 encoder 主干用较小的学习率比如 5e-5。原因是分类头是随机初始化的需要学得更多主干是从预训练来的只需要微调。用分组学习率同时照顾两边的需求微调初期的 loss 下降明显更平滑。第二个技巧是关于输入 patch size 的选择。ViT 的 patch size 可以从 16 改成 8代价是 token 数量变成原来的 4 倍计算量显著上升但在细粒度分类任务上准确率确实能提升。如果你做的是花卉、昆虫、车辆型号这类差距很小的分类任务可以试试 patch16 预训练 微调时改用较小 patch size 的做法。不过要记住改小 patch size 之后position embedding 的插值不可避免微调时最好把位置编码也一起微调否则位置信息会不匹配。ViT 这个架构从出现到现在已经过了好几年中间又演化出了 Swin、DeiT、PVT 等一大堆变体但它的核心思路——把图像变成序列、用注意力建模全局关系——至今仍然是视觉模型的主流方向之一。把 ViT 本身吃透后面再接触任何 Transformer-based 的视觉模型都会轻松很多。希望这篇拆解能帮你省掉我当年摸索的弯路。