最近又把 ViT 这篇《An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale》翻出来读了一遍。原因很简单这两年 Transformer 在视觉领域已经快成默认起点了从检测到分割再到视频理解到处都能看到它的影子但很多朋友讨论问题的时候其实对原始论文里的设计动机并不清楚。我自己做图像分类项目从 CNN 迁到 Vision Transformer 也踩过不少坑回头重读原文才发现不少细节论文里其实早就有提示只是当初读得太快没当回事。这篇笔记想把论文从动机、架构、实验到复现踩坑完整过一遍。如果你正准备第一次读 ViT或者已经用过 ViT 但不清楚里面每个部件为什么存在、训练时该怎么设参数那这篇文章应该能帮你省下不少时间。我不打算做逐段的论文翻译而是把关键设计拆开讲透再加上实操层面的经验争取让你读完既能理解论文也知道怎么做实验。1. 论文到底在解决什么问题1.1 为什么非要用 Transformer 处理图像ViT 论文的核心主张用作者的话来说就是图像识别并不一定要依赖 CNN 里的那些卷积操作一个纯粹基于自注意力的架构只要给它足够多的数据去训练就能在 ImageNet 这种大规模基准上达到甚至超过 CNN 的效果。这句话放到今天看好像没什么稀奇Transformer 都在做各种视觉任务了但你要回到 2020 年当时的主流共识还是卷积天然适合图像Transformer 只是 NLP 的玩具。CNN 的成功建立在两个归纳偏置上局部性和平移等变性。卷积核只感知局部区域通过堆叠层来扩大感受野同一套卷积核在图像不同位置共享所以猫在图片左上角还是右下角都能被认出。这种先验让 CNN 在数据量有限时也能学得很好但反过来也限制了它的上限——它天生就看不到全局关系或者说得更准确一点要看全局需要把感受野堆得足够大效率很不划算。Transformer 几乎没有这种归纳偏置。自注意力机制随便来两个位置都能直接计算依赖关系第一层就可以看到整张图它也不假设位置平移等变所以需要一个额外的位置编码把像素在哪告诉模型。用大白话讲CNN 像一个靠本地邻居消息传话的网络而 ViT 直接把所有人拉进一个会议室让信息自由流动。这样更灵活但代价就是模型得从纯数据里把图像结构重新学一遍这也是为什么 ViT 在大数据上能超过 CNN在中小数据上反而打不过。1.2 从图像到词序列patch 的直觉NLP 里的 Transformer 拿到的输入是一串 token每个 token 是文本里的一个词或者子词。图像怎么往这个框架里套这是整篇论文设计中最关键的一次决策。作者尝试过几乎像素级别地把每个像素当作一个 token但这有两个不可接受的问题。一是序列长度爆炸一张 224x224 的图就有 50176 个像素注意力复杂度是序列长度的平方根本算不动。二是单个像素的语义太弱一个像素本身学不到什么有意义的特征靠它去构建全局关系缺乏中间层的结构支撑。所以作者提出了 patch embedding 的思路把一张图切成一个个 16x16 的小块每个小块展平后做一次线性映射就得到一个 token。16x16 这个数字不是随便选的。在 224 分辨率下16x16 的小块恰好生成 196 个 token224/16 等于 1414x14 等于 196和 NLP 场景里常见的序列长度在数量级上对得上。同时16x16 的 patch 已经包含了足够的局部纹理和结构信息相当于把像素级的原始输入提升到了语义颗粒度的输入模型第一层就能在局部结构和全局关系之间找一个平衡点。这个思路本质上和我们看马赛克拼图很像每一块虽然只是局部图案但拼到一起能看出完整的画面而且你一眼扫过去最先建立的是整体轮廓而非单块细节。1.3 ViT 与 CNN 的本质差异接下来要分清 ViT 和 CNN 到底差在哪儿这个理解决定了你训练和使用它们的方式。CNN 的感受野是从小到大逐层扩大的浅层看边角深层看形状天然符合由局部到整体的视觉处理习惯。但 CNN 对这种先验的依赖很强一旦某个特征超出了现有感受野就要继续堆层运算量很快就上去了。ViT 每一层都是全局交互。任意两个 patch 之间可以直接建立依赖不管它们在图上隔多远。等于第一层就已经把全局信息和局部信息混合在一起后面层继续提炼关系。这既是优势也是劣势——没有 CNN 那种平滑的渐进式抽象所以它对数据量特别贪婪。CNN 的平移等变性让它在小数据上表现稳定ViT 没有这个免费午餐所以论文里反复强调要在超大规模数据上预训练。数据量不够时ViT 学到的关系实际上是偷懒的它可能硬记住了一些训练集特有的伪模式泛化性自然就差了。我用一个项目举例。之前做一个商品图像分类任务训练集只有不到十万张图直接从头训 ViT结果验证集准确率比同条件下的 ResNet 低了大概三个点而且 loss 波动明显。后来换成熟预训练权重做微调同一个 ViT 反而比 ResNet 高了两个点。这个体验和论文里的表 1 结论完全一致性能表现跟随预训练数据规模的曲线走数据少CNN 赢数据多Transformer 赢。2. 关键架构细节拆解2.1 Patch Embedding 的实现与数学本质假设输入图像为 x尺寸是 H×W×C。ViT 的 patch embedding 过程可以分三步走。第一步把 H×W 的平面切成若干 P×P 的小块。如果 H 和 W 都等于 224P 等于 16那么每一维上有 14 个小块总块数 N 等于 14×14196。每个小块原始维度是 16×16×3768。第二步把每个小块展平成向量。展平后的维度是 768对应每个 patch 的所有像素值排在一起。第三步用同一个可学习的线性映射把 768 维向量投影到模型宽度 D 维度。论文默认设置 D 等于 768这一步等于一个全连接层输入维度 768输出维度 768。如果你选更大的模型比如 ViT-Large那 D 可能是 1024这一步的输出维度就变成 1024。公式层面等价于 x_patch 乘以一个形状为 (P²C, D) 的权重矩阵再加一个可学习的偏置。在 PyTorch 里最干净的做法是用一个 Conv2d 层来实现import torch import torch.nn as nn class PatchEmbed(nn.Module): def __init__(self, in_channels3, patch_size16, embed_dim768): super().__init__() self.patch_size patch_size # 卷积核大小和步长都设为 patch_size效果等价于切 patch 线性投影 self.proj nn.Conv2d(in_channels, embed_dim, kernel_sizepatch_size, stridepatch_size) def forward(self, x): # x: (B, 3, H, W) - (B, embed_dim, H/P, W/P) - (B, N, embed_dim) x self.proj(x) x x.flatten(2).transpose(1, 2) return x这里用卷积一步实现切块和线性投影是因为当一个卷积核的尺寸和步长恰好等于 patch_size且输入通道到输出通道的映射是线性时它所做的运算和展平 patch 再过全连接是完全等价的。理解这一点对后续做分辨率适配很有帮助因为改输入分辨率时只需要保证 H 和 W 都能被 patch_size 整除就行。2.2 位置编码与 CLS token 的作用Transformer 本身对输入的序列顺序不敏感打乱 patch 的顺序之后自注意力计算出来的结果完全一样。但图像的空间信息显然是重要的——上面和下面对模型的意义完全不同。所以 ViT 加入了一个可学习的 1D 位置编码形状为 (N1, D)这里的 N1 是 196 个 patch token 加 1 个分类 token。论文里对比过几种位置编码方案可学习的 1D 编码、可学习的 2D 编码、以及正弦余弦的 2D 编码结果发现它们在最终精度上几乎没有差别。这一条实验结果经常被人忽略但它很有意思这说明模型其实可以靠内容本身判断两个 token 的位置关系并不需要刻意把位置信息做得特别精细。就像你在一张照片中看到天空在图片上方、地面在下方即使不给你坐标信息你也能根据内容推断出相对位置。所以作者最终选了最简单的可学习 1D 编码避免引入不必要的复杂度。CLS token 的设计则直接借用了 BERT。在输入序列的最前面放一个特殊的 token它在所有层里都参与注意力计算最终分类的时候只取这个 token 对应的输出向量过一层 LayerNorm 再接一个线性分类头。为什么不把 196 个 patch token 的输出做平均或最大池化作者在实验里试过用 CLS token 和用全局平均池化效果都非常接近但 CLS token 有一个好处它的角色一目了然并且后续如果要做一些跨模态或 dual-encoder 的项目CLS token 直接充当全局表征用起来很方便。我自己做图文匹配的时候就是直接取 CLS token 当特征向量去算相似度省去了池化的步骤。2.3 多头自注意力的计算流程自注意力是 ViT 每个 Block 的核心。设输入序列为 Z形状是 (N1, D)。首先通过三个独立的线性层计算出 Q、K、V形状都是 (N1, D)# 伪代码仅用于展示计算逻辑 Q Z W_q # (N1, D) (D, D) - (N1, D) K Z W_k V Z W_v接着把 Q、K、V 按头数 h 拆开每个头的维度变成 D/h。对每个头独立计算注意力权重# 每个头里计算 scores Q_head K_head.T / sqrt(D/h) # 缩放点积注意力 weights softmax(scores, dim-1) output_head weights V_head最后把所有头的输出拼接回去再过一层输出线性投影。这就是 Multi-Head Self-Attention 的完整流程。缩放因子 sqrt(D/h) 是为了防止点积结果过大导致 softmax 梯度过小这一点在 Transformer 原始论文里就解释过ViT 原封不动搬了过来。每个 Block 除了注意力还有一个 MLP 模块由两个全连接层夹一个 GELU 激活组成。第一个全连接把维度从 D 扩到 4D第二个再缩回 D。这个 4 倍扩展比例也是经典配置实践证明在宽度和容量之间取了一个不错的平衡。每个 Block 都用了 Residual Connection 和 LayerNorm顺序是 LayerNorm 在前、注意力或 MLP 在后也就是 Pre-LN 结构。Pre-LN 的好处是训练更稳定能支持更大的学习率论文里从头训练 ViT 能训到 300 个 epoch 不炸和这个结构选择密切相关。3. 实验设计与关键结论3.1 预训练数据规模如何改变结论论文里最著名的一组实验是拿 ViT 在三个不同规模的数据集上训练再对比 BiTBig Transfer一个大规模预训练残差网络和基于 ResNet 的其他变体。在中等规模数据集相当于百万级图片的公开基准上ViT 只能拿到中上成绩和同容量的 ResNet 相比并没有优势甚至某些设置下略差。这个现象的原因前面提过Transformer 没有空间归纳偏置数据量不足以让模型从零学会图像结构它就倾向于欠拟合或者过拟合。但当预训练数据规模提升到数亿张图片级别时情况彻底反转。ViT 的准确率超过了同等甚至更大规模的 CNN 模型而且训练所需的算力比 CNN 更少。作者在论文里给了非常直白的解释卷积等价于一种受限的自注意力它只在一个局部窗口内计算而且权重共享方式固定如果让 Transformer 在大数据上自由学习它可以自己学会卷积的行为也可以发现比卷积更好的交互方式。这个结论对实际项目有直接指导意义。我自己之前遇到过数据量不够但想强上 Transformer 的场景正确做法不是改模型结构而是先问自己能不能找到一个足够大的预训练数据集或者干脆用公开的大规模预训练权重做初始化。否则你亲手训出来的 ViT 很可能让你怀疑人生。可以说预训练规模这个变量比 ViT 模型本身任何超参数都重要。3.2 微调阶段的分辨率调整与位置编码插值ViT 的标准做法是先在 224×224 分辨率下预训练然后在微调阶段把输入分辨率提高到 384 或 512。这里有一个细节特别值得注意。当分辨率提高时patch 大小保持 16×16 不变所以 patch 数量会变多。比如从 224 提到 384patch 数量就从 196 变成 576。预训练时的位置编码长度为 197196 加 CLS现在序列长度变成 577直接对不上。ViT 的解决办法是对预训练的位置编码做 2D 插值再配合微调训练让模型适应新分辨率。我第一次复现的时候踩过一个坑想当然地对位置编码矩阵做简单的一维线性插值结果准确率掉了好几分。后来按论文的方式先把它 reshape 成 14×14 的 2D 形态再双线性插值到 24×24最后展平效果才正常。为什么这样对了因为图像的空间关系是二维的对行方向和列方向分别插值符合图像的几何直觉直接按一维插值会把相邻行之间的距离关系搞乱模型丢失了原有的位置先验。微调的时候还有一个细节分辨率提高后序列里 patch token 之间的交互距离变了但实际上只需要微调十几到几十个 epoch 就能收敛得不错因为初始化权重已经包含了很强的特征表达能力。论文里给的微调学习率远小于从头训练的学习率同时禁用了一些强正则化比如把权重衰减调低或者不用 Mixup以免破坏已经学好的特征。3.3 不同模型规模下的表现对比ViT 论文定义了几个标准模型规格ViT-BaseD76812 个 Block12 个头参数约 86M、ViT-LargeD102424 个 Block16 个头参数约 307M、ViT-HugeD128032 个 Block16 个头参数约 632M。论文用 JFT 这种数亿级数据集预训练之后在 ImageNet 上微调ViT-Large 已经超过当时最好的 BiT 和 EfficientNet 变体ViT-Huge 则更进一步。一个容易忽视的点是大模型的训练成本并不像参数数量看上去那么吓人。作者在论文中对比过训练吞吐量同等精度下ViT-Large 比 ResNet 系列的训练成本低了不少因为 GPU 对矩阵乘法特别擅长Transformer 的全局结构让它在张量核心上利用得非常充分。这一点对工程选型很有参考价值同样算力预算下与其堆一个超深的卷积网络不如试试用张量并行跑 ViT吞吐量和效果可能都更好。4. 实操复现配置、代码与避坑4.1 推荐训练配置与超参数选择从零训练一个小规模 ViT 用来做实验我建议按下面这套配置开始它在稳定性和效果之间比较平衡# 优化器AdamW # 学习率0.001使用 10 个 epoch 线性 warmup # 权重衰减0.05 # batch size1024如果显存不够就减小同步调低学习率 # 训练轮数300 epoch # 数据增强Rand-Augment Mixup CutMix # 标签平滑0.1 # 随机深度0.1ViT 在小数据集上从头训练经常不收敛我把最大头的两个原因列一下Batch size 太小。ViT 对 batch size 比较敏感batch 太小的时候BatchNorm 或 LayerNorm 统计不稳定容易过拟合。实测在 8 卡 GPU 上batch size 低于 256ViT-Base 的表现就会下滑明显。学习率太高且没有 warmup。Transformer 类模型前期梯度震荡剧烈如果没有 warmup 直接跑大学习率loss 很容易冲到 NaN。我一般至少留 5 到 10 个 epoch 做线性 warmuplearning rate 从 0 平缓升到目标值。4.2 一个可用的 PyTorch 核心代码下面是一份精简但可以训起来的 ViT 核心结构适合二次魔改。我把分类头单独抽出来方便你自己适配不同任务。import torch import torch.nn as nn from einops import rearrange class Attention(nn.Module): def __init__(self, dim, heads12, dropout0.): super().__init__() self.heads heads self.scale (dim // heads) ** -0.5 self.qkv nn.Linear(dim, dim * 3) self.attn_drop nn.Dropout(dropout) self.proj nn.Linear(dim, dim) self.proj_drop nn.Dropout(dropout) def forward(self, x): B, N, C x.shape qkv self.qkv(x).reshape(B, N, 3, self.heads, C // self.heads) q, k, v qkv.permute(2, 0, 3, 1, 4) 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_drop(self.proj(x)) class Block(nn.Module): def __init__(self, dim, heads12, mlp_ratio4., dropout0.): super().__init__() self.norm1 nn.LayerNorm(dim) self.attn Attention(dim, heads, dropout) self.norm2 nn.LayerNorm(dim) self.mlp nn.Sequential( nn.Linear(dim, int(dim * mlp_ratio)), nn.GELU(), nn.Dropout(dropout), nn.Linear(int(dim * mlp_ratio), dim), nn.Dropout(dropout), ) def forward(self, x): x x self.attn(self.norm1(x)) x x self.mlp(self.norm2(x)) return x class ViT(nn.Module): def __init__(self, img_size224, patch_size16, in_channels3, num_classes1000, dim768, depth12, heads12, dropout0.): super().__init__() self.patch_size patch_size self.num_patches (img_size // patch_size) ** 2 self.patch_embed nn.Conv2d(in_channels, dim, kernel_sizepatch_size, stridepatch_size) self.cls_token nn.Parameter(torch.zeros(1, 1, dim)) self.pos_embed nn.Parameter(torch.zeros(1, self.num_patches 1, dim)) self.pos_drop nn.Dropout(dropout) self.blocks nn.Sequential(*[ Block(dim, heads, dropoutdropout) for _ in range(depth) ]) self.norm nn.LayerNorm(dim) self.head nn.Linear(dim, num_classes) def forward(self, x): B x.shape[0] x self.patch_embed(x) # (B, dim, H/P, W/P) x rearrange(x, B D H W - B (H W) D) 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[:, 0]) # 取 CLS token 输出 return self.head(x)这份代码的 patch embedding 用 Conv2d 实现后面直接接 Transformer Block。注意init_weights部分没有写实操里我习惯用截断正态分布初始化所有线性层并把 pos_embed 和 cls_token variance 设成约 0.02否则训练前期 attention 输出方差会太大导致 loss 下降奇慢。4.3 ViT 比 CNN 更依赖正则化与数据增强前文提到ViT 缺少卷积的归纳偏置所以它天生需要更强的正则化手段来防止过拟合。论文里的实验设置其实已经给出了答案Rand-Augment、Mixup、CutMix、标签平滑、DropPath 一起上。这些手段具体加多少我建议参考原论文的做法中等数据规模下Mixup alpha 设为 0.8CutMix alpha 设为 1.0Rand-Augment 的幅度设为 9标签平滑系数 0.1。如果你要在自定义数据集上复现先照抄这个配置再逐步调整比从头猜参数靠谱得多。DropPath 在 ViT 里是一个很有用的正则化工具它按概率在训练时随机丢弃整个残差分支等效于在深度方向上做集成。对于深层 ViTDropPath rate 设 0.1 到 0.3 对最终精度有明显帮助。很多蒸馏类 ViT 模型的技术报告里也都保留了这个设置可见它是普遍有效的。5. 常见问题与排查记录5.1 小数据集上训练 ViT 不收敛怎么办这个问题我几乎每周都能在技术社区看到人问。在几千张图片的数据集上从头训 ViT结果 loss 下降很慢或者干脆发散。这属于正常现象而不是代码写错了。我的建议按优先级排列换用预训练权重做初始化然后在目标数据上微调而不是从头训练。如果用不了预训练权重那就别再自己搭深层 Transformer试试浅层加宽的结构比如 6 个 Block、宽度 512再配强数据增强。把学习率降到 3e-4 以下用 large batch 加 warmup先把 loss 训到一个正常的下降曲线上再考虑提升精度。我踩过的具体例子是用两万张图从头训 ViT-Small前期 loss 一直在 3 到 4 之间徘徊降到 0.001 并给足 warmup 之后曲线才开始正常下降。ViT 对学习率的敏感程度远超 CNN这算是它最反直觉的一点。5.2 位置编码插值后模型掉点上述问题通常出现在用 ImageNet 预训练权重迁移到更高分辨率输入时。此时位置编码的序列长度不匹配需要对 pos_embed 做插值。正确的操作分成两步把 pos_embed 中的 CLS token 分离出来只对剩下的 patch token 位置编码进行插值。把 patch token 位置编码 reshape 成二维网格比如从 (14, 14) 插值到 (24, 24)使用双线性插值再和 CLS token 位置编码拼接回去。注意插值之后要接着微调一定轮数让模型适应新的位置编码分布。不能插值完就不训练直接评估那样位置信息对不上模型精度一定受损。我自己实验里微调 20 个 epoch 内插值插得好不好能拉开两个点左右。5.3 显存占用太大训练跑不起来ViT 在大分辨率输入下序列长度会变长注意力矩阵的显存占用和数据规模成正比。比如输入分辨率从 224 提到 384序列长度从 196 变成 576注意力矩阵大小大约是原来的 3 倍显存压力非常大。对应解决办法有三个用 FlashAttention 这类融合算子它不显式构造完整注意力矩阵而是分块计算能把显存占用降一个量级。开启梯度检查点gradient checkpointing相当于用少量计算换显存forward 的时候不保存中间激活backward 时重新算一遍。减小 batch size 或者用混合精度训练这也是最直接的降显存手段。在推理阶段如果显存也紧张可以试试用一个较小的 patch size 替代大 patch size但要注意这样会改变输入序列长度和预训练权重的配置不完全兼容必要时还得配合位置编码插值。5.4 ViT 推理速度是否真的比 CNN 慢这个问题要分开看。ViT 的全局自注意力在序列长度短的时候计算量不算夸张但 patch 数量多起来之后复杂度呈平方增长。所以在 224 输入、patch 16 的情况下ViT-Base 的推理速度和同量级的 ResNet-50 相比不相上下而且 GPU 上的吞吐量甚至可能更高。但到了 384 甚至 512 分辨率下ViT 的推理成本会明显高于同精度 CNN。工程上如果要在边缘设备上部署最好先考虑蒸馏方案用 ViT 做教师模型去蒸馏一个轻量 CNN 或小型 ViT精度损失可控速度和部署友好度大幅提升。这个方案我实际做过效果不错教师的全局建模能力确实能迁移到学生模型上。6. 论文之外的个人体会读 ViT 论文你会发现它最大的价值不只是提出了一个新架构而是做了一个非常清晰的实验论证归纳偏置不是免费的也不是越多越好。CNN 内置的局部性假设在小数据上很管用但在大数据场景下反而可能限制了模型的上限。ViT 敢于去掉这些先验然后用数据和算力去弥补最终在超大规模预训练下实现了反超。这个思路对后来很多视觉架构的设计产生了深远影响也让我在做模型选型时多了一个判断维度你的数据量到底撑不撑得起一个少归纳偏置的模型。最后分享一个小技巧如果你项目里数据量不多但就是想用 Transformer 结构可以先别急着从头训也别急着改网络结构。先把数据增强策略堆足再去一个大模型做蒸馏或者直接用预训练权重微调。我实验多次发现这一套组合拳通常比你在结构上折腾来折腾去更有效。ViT 的很多成功实践说白了都是围绕数据喂够了吗和正则化做足了吗这两个问题展开的。抓住这两点你的 Transformer 视觉项目就成功了一大半。