Vision Transformer通常简称ViT可以说是我近几年研究视觉模型时开销最大、也最值得的一次思维重置。过去我们默认卷积神经网络就是图像特征提取的万能解直到ViT出现把NLP领域的Transformer搬到了图像里而且是用最直接、最暴力的方式把图片切成一个个小方块当成一个序列喂给注意力机制。结果大家都知道了在没有卷积的前提下它在大规模数据上的表现竟然全面超越了当时的CNN模型这让我当时看的目瞪口呆。对于正在入门视觉方向的开发者或者已经在用CNN做项目但想了解下一代架构的工程师这篇文章会从原理、代码、训练技巧、踩坑实录四个方面把Vision Transformer拆开揉碎地讲清楚。我会用我自己复现和调优的实战经验说话不堆公式、不绕概念保证每一步你都能在本地跑起来。1. 为什么需要Vision Transformer从CNN的局限说起很多人一上来就研究ViT的代码但我建议你先搞清楚一个问题CNN为什么不够用了只有理解了动机你才能明白ViT里那些看似奇怪的设计到底为什么存在。1.1 CNN的感受野与长距离依赖问题CNN的卷积核比如3×3或7×7天然只能感知局部区域。虽然通过堆叠层数高层的特征图能间接覆盖到较大的区域但这种长距离依赖是非常间接且脆弱的。简单说CNN处理一张图时左上角的猫头和右下角的猫尾巴如果要建立关联需要经过很多层卷积的逐步传递信息在传递过程中可能被稀释或丢失。尤其在一开始的前几层每个位置只能看到它周围那一小块如果你要识别一只猫的整体轮廓卷积层需要慢慢拼图。这个过程中误判的概率很高而且对全局关系的建模非常不直接。我在某些场景里用CNN做小目标检测时发现模型总是抓不住物体之间的语义关联比如人骑着自行车这种关系靠卷积网络硬学真的很吃力。它更多是在学局部的纹理和形状而不是全局的布局和关系。1.2 自注意力机制的核心思想自注意力机制的思路恰恰相反它在处理任何一个位置时会直接计算它跟图片上所有其他位置的相关性然后在相关性的指导下聚合信息。这就像是一张照片里每个像素都能直接跟其他所有像素对话不需要邻居转达。注意力矩阵告诉你这个区域跟哪个区域关系最密切然后根据密切程度加权求和。这个机制直接解决了长距离依赖问题。在图像里长距离可能意味着图片对角线上的两个目标也可能意味着完全不同的语义区域。自注意力没有距离概念只要位置在序列里不管隔多远都能一步到位地建立关联。这就是为什么Transformer在语言领域那么强。1.3 ViT的出场把图像当作句子处理ViT的绝妙之处在于它把图像变成了句子。具体来说把一张图分成固定大小的patch每个patch展平成向量就像句子里的一个词patch的空间位置就像词在句子里的位置。然后整个序列直接丢进标准Transformer编码器里。这个转换干净利落没有引入复杂的特殊结构完全复用NLP那一套。我当时看到这个设计的第一反应是这得吃多少计算量确实ViT的信息密度比CNN低需要更多数据才能训练好。原论文也是在大规模数据集上做了预训练然后在中等数据集上微调才展示出优势。但这并不妨碍它的理念成为主流趋势。从那以后视觉领域的设计逻辑几乎被重写除非你是要做轻量级、实时推理的端侧应用否则一谈起通用视觉架构大家第一反应就是ViT或者它的变体。2. 核心原理拆解ViT是如何工作的下面我们从输入到输出把ViT的每个部件逐一拆开。这部分建议你边读边对照代码理解会深入很多。2.1 图像切块与嵌入patch嵌入的完整流程假设输入图像是224×224×3我们设定patch大小为16×16。那么图像会被切成 (224/16)^2 14×14 196个patch每个patch的维度是 16×16×3 768。这一步在代码里通常用一层的卷积来完成卷积核大小等于patch大小步长等于patch大小。用卷积的好处是能利用GPU加速而效果跟手动切割完全一致。随后每个patch会被展平成768维的向量。这个向量是原始像素不能直接进Transformer需要用一层全连接把它线性投影到一个可学习的嵌入空间。这个投影层的输出维度通常记为D也就是Transformer隐藏层维度比如768或1024。原始论文里也加了一个类别标记class token它初始化为一个可学习的向量跟图像patch拼在一起位置放在序列的最前面。这个class token在Transformer的编码过程中充分聚合了全局信息最后拿出来接分类头相当于整张图的整体表征。2.2 位置编码给patch加上空间坐标Transformer本身没有顺序概念所以我们必须把patch的位置信息显式加进去。这里有两种主流做法绝对位置编码和相对位置编码。ViT开山之作用的是绝对位置编码它是一个可学习的参数矩阵形状为 (num_patches 1, D)其中1就是class token。在训练过程中这个位置向量会随着网络一起更新。有个细节让我花了不少时间位置编码对图像平移是敏感的。也就是说如果图片内容整体移动了几个像素patch的边界可能完全改变导致同样的物体落在不同的patch组合里。这就削弱了模型的平移等变性而这是CNN天然具备的。所以ViT一般需要更充足的数据增强或者更大的数据集来弥补这一点。在实践里如果你的数据量不够位置编码可能学到的是训练集中常见的空间分布泛化性会受影响。后面我会讲怎么减轻这个问题。2.3 Transformer编码器多头自注意力与MLP把patch嵌入和位置编码相加后得到序列长度等于1961197、维度为D的矩阵。接着进入Transformer编码器它由若干个相同的层堆叠而成。每一层里有两个核心子层多头自注意力MSA和前馈网络MLP。两个子层外面都有残差连接和层归一化标准流程是先做层归一化再进子层之后加回原输入最后再走一次归一化和MLP。多头自注意力简单理解就是让模型同时从多个不同的关系角度去扫描图像。比如一个头关注颜色相似区域的关联另一个头关注形状边缘的关联。每个头的维度是D/hh是头的数量。计算流程大致是每个token会生成三个向量——Query查询、Key键、Value值然后每个token的Query去跟所有token的Key做点积经过softmax得到注意力权重再用权重对Value加权求和。这个步骤扩展开来就是矩阵乘法。其中计算量的大头在注意力矩阵序列长度197、维度D的序列计算出的注意力矩阵是197×197这个矩阵的规模随序列长度平方增长。MLP部分则是两个全连接层中间往往有GELU激活函数和Dropout。MLP对整个模型的参数贡献最大通常基座的隐藏维度要放大到4倍。ViT的核心思想就是重复堆叠这种结构层数越多表达的抽象层次越高但训练难度也相应增加。我一般用的配置是12层D768头数12参数总量在86M左右。这个量级在不算大的数据集上已经能学到非常鲁棒的特征。2.4 分类头与训练目标在序列的最后一个输出处对应class token的那个位置接一个全连接分类头输出类别数。整个模型用标准的交叉熵损失进行端到端的训练。值得注意的是在预训练阶段还会训练一个辅助的解码器或者额外的对比学习损失但主流的ViT训练通常就是在分类目标上直接进行。后面只要用分类头微调就行。这里有一个容易误解的点class token和全局平均池化相比哪个更好我在实践里发现class token的意义在于它本身就是一个可学习的线性探针它在编码过程中会根据数据自动决定该从所有patch里提取什么信息。而如果用平均池化等于是告诉模型每个位置都平等贡献。两者在性能上有差异但通常不大。关键是代码实现时不要搞混尤其是把class token从整个序列中拆分出去再接分类头的步骤。3. 实操过程从零实现一个Vision Transformer接下来是大家都关心的部分怎么写代码、怎么跑起来。我会以PyTorch为例手把手搭建一个最小可用的ViT然后说明训练细节。这里用的是一个虚构的模拟数据集你可以替换成自己的数据。3.1 环境准备与数据说明你需要Python 3.8以上PyTorch 1.10以上以及GPU最好。如果只有CPU可以把模型调小或者使用小patch。我的建议是使用官方ImageNet格式的文件夹数据或者直接用torchvision里的数据集做测试。为了快速验证模型结构是否正确可以先创建一个小批量的随机张量测试前向传播再正式训练。pip install torch torchvision einopseinops是一个非常优雅的张量重排库我们用它来实现patch嵌入会极其简洁。3.2 模型结构代码从PatchEmbed到完整ViT下面我直接给出一个简洁版的ViT实现。这段代码经过了多次简化但保留了核心逻辑。你不需要用工具库可以直接复制运行。import torch import torch.nn as nn from einops import rearrange class PatchEmbed(nn.Module): def __init__(self, in_channels3, patch_size16, embed_dim768): super().__init__() self.patch_size patch_size self.proj nn.Conv2d(in_channels, embed_dim, kernel_sizepatch_size, stridepatch_size) def forward(self, x): # x: (B, C, H, W) x self.proj(x) # (B, embed_dim, H/p, W/p) x rearrange(x, b d h w - b (h w) d) # (B, num_patches, embed_dim) return x class MultiHeadSelfAttention(nn.Module): def __init__(self, embed_dim, num_heads): super().__init__() assert embed_dim % num_heads 0 self.dim_head embed_dim // num_heads self.num_heads num_heads self.qkv nn.Linear(embed_dim, embed_dim*3) self.proj nn.Linear(embed_dim, embed_dim) self.scale self.dim_head ** -0.5 def forward(self, x): B, N, D x.shape qkv self.qkv(x).reshape(B, N, 3, self.num_heads, self.dim_head) 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) out (attn v).transpose(1,2).reshape(B, N, D) return self.proj(out) class MLP(nn.Module): def __init__(self, embed_dim, mlp_ratio4.0, dropout0.): super().__init__() hidden_dim int(embed_dim * mlp_ratio) self.fc1 nn.Linear(embed_dim, hidden_dim) self.act nn.GELU() self.fc2 nn.Linear(hidden_dim, embed_dim) self.dropout nn.Dropout(dropout) def forward(self, x): x self.fc1(x) x self.act(x) x self.fc2(x) return self.dropout(x) class TransformerBlock(nn.Module): def __init__(self, embed_dim, num_heads, mlp_ratio4.0, dropout0.): super().__init__() self.norm1 nn.LayerNorm(embed_dim) self.attn MultiHeadSelfAttention(embed_dim, num_heads) self.norm2 nn.LayerNorm(embed_dim) self.mlp MLP(embed_dim, mlp_ratio, dropout) self.dropout nn.Dropout(dropout) def forward(self, x): x x self.dropout(self.attn(self.norm1(x))) x x self.dropout(self.mlp(self.norm2(x))) return x class VisionTransformer(nn.Module): def __init__(self, image_size224, patch_size16, in_channels3, num_classes10, embed_dim768, depth12, num_heads12, mlp_ratio4.0, dropout0.): super().__init__() num_patches (image_size // patch_size) ** 2 self.patch_embed PatchEmbed(in_channels, patch_size, embed_dim) self.cls_token nn.Parameter(torch.zeros(1, 1, embed_dim)) self.pos_embed nn.Parameter(torch.zeros(1, num_patches1, embed_dim)) self.pos_drop nn.Dropout(dropout) self.blocks nn.Sequential(*[ TransformerBlock(embed_dim, num_heads, mlp_ratio, dropout) for _ in range(depth) ]) self.norm nn.LayerNorm(embed_dim) self.head nn.Linear(embed_dim, num_classes) nn.init.trunc_normal_(self.cls_token, std0.02) nn.init.trunc_normal_(self.pos_embed, std0.02) nn.init.kaiming_normal_(self.patch_embed.proj.weight) def forward(self, x): B x.shape[0] x self.patch_embed(x) cls_token self.cls_token.expand(B, -1, -1) x torch.cat([cls_token, x], dim1) x x self.pos_embed x self.pos_drop(x) x self.blocks(x) x self.norm(x) cls_out x[:, 0] return self.head(cls_out)上面这段代码足够跑起来。每次先实例化再构造一个随机张量试一下输出维度。model VisionTransformer(image_size224, patch_size16, num_classes1000) fake torch.randn(1, 3, 224, 224) out model(fake) print(out.shape) # torch.Size([1, 1000])3.3 训练参数设置与计算量分析训练ViT的时候最怕的是参数设置不对直接导致训练崩溃。从我复现的经验看有几个关键参数要特别留意。首先学习率不能像CNN那样给太大一般基础学习率设为0.001配合warmup和cosine decay。warmup是必要的因为Transformer的结构较深一开始就大学习率会让层归一化和注意力权重剧烈震荡。warmup步数一般是总的训练步数的5%到10%。batch size也很关键我建议尽可能用大batch比如128或以上。但是ViT结构本身显存占用大如果你显卡不够可以用混合精度训练、梯度累积。我的经验是batch size减小后要适当提高学习率不对实际上batch小学习率也要小。一般线性缩放规则就是batch size翻倍学习率也翻倍。optimizer torch.optim.AdamW(model.parameters(), lr1e-3, weight_decay0.05) scheduler torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_maxepochs, eta_min1e-5)还有一个很多人忽略的问题初始化。class token和position embedding我用的是trunc_normal初始化。如果全都初始化为0注意力分数一开始可能过于均匀导致所有位置的信息混合。实践上标准初始化配合小的std0.02可以加快收敛。3.4 训练过程与结果观察在模拟数据上我一般先跑个10个epoch观察loss下降曲线。如果loss没有下降检查是否数据归一化有问题。ViT对输入分布比较敏感一定用标准化的图像。另外因为模型参数大极易过拟合所以增强策略一定要重一些随机裁剪、翻转、RandAugment都是必备。我在一个中等规模的数据集几万张图上训练时发现ViT的收敛速度明显慢于ResNet但后期损失会降得更低。这个规律印证了它需要更多数据如果数据不够很容易卡在一个较高的loss上而且验证集准确率波动大。4. 训练ViT的实战经验与关键技巧模型能跑通是一回事要真正用ViT做项目必须掌握下面这些实战经验否则你可能永远停留在代码能跑效果不行的阶段。4.1 数据量不足怎么办预训练与数据增强ViT最被人诟病的一点就是数据饥饿。在ImageNet量级的百万张图片上ViT需要至少300M的预训练数据才能超过CNN。但这不代表小数据集上完全不能玩。关键在于三种手段一是使用预训练权重。从公开的ImageNet预训练模型初始化然后在小数据集上微调这是最现实的做法。比如你用torchvision或timm库加载一个现成的ViT权重然后替换分类头冻结前几层只训练最后几层和分类头。我试过在只有一两万张的数据集上微调轮次拉长效果比从头训CNN好很多。二是数据增强。RandAugment、MixUp、CutMix等增强方法对ViT特别重要。因为ViT缺乏CNN的归纳偏置它需要靠增强来认识到物体旋转和平移都是同一个物体。我给ViT组做实验时发现使用CutMix的模型在验证集上能稳涨3到5个点而CNN只涨1到2个点。这说明增强对ViT的价值远大于CNN。三是知识蒸馏。如果你想自己从零训一个轻量ViT可以先用一个大的CNN模型生成软标签然后让ViT去学软化后的概率分布。这个操作能显著提升小模型的准确率本质上相当于把CNN的归纳偏置蒸馏给ViT。4.2 优化器、学习率与正则化的正确组合ViT标准的优化配置是AdamWweight_decay给到0.05甚至0.1。为什么这么高因为ViT的参数量大如果没有强大的权重衰减很容易记住训练集的噪声。还有一个细节图片增广和随机噪声都会带来loss震荡但不要因此调低weight decay反正我在实验中发现高weight decay反而加快收敛因为它等于一直做某种平滑。学习率的调度必须是warmupcosine decay这个组合几乎成为ViT训练的黄金公式。warmup让训练初期注意力矩阵不至于被极端梯度破坏cosine decay让后期学习率缓慢下降有助于微调权重。另外回头说一句dropout可以加在Attention层和MLP层但不要加在block的加和之后。如果你发现验证集与训练集loss差距大优先提高dropout而不是降低模型复杂度因为ViT本身设计就是大而全剪枝会导致它丧失注意力优势。4.3 位置编码的泛化一个容易被忽略的坑ViT在预测时图像尺寸必须和训练时一致或者至少是整数倍关系。原因是位置编码矩阵的形状是固定的。如果换了更大的分辨率比如把256×256的图像输入给224×224训练的模型position embedding数量对不上直接报错。解决办法有几种第一图像缩放回训练尺寸这是最简单但会损失信息第二使用插值把位置编码扩展成新的尺寸但插值会破坏编码的相对位置语义通常还需要微调第三换成相对位置编码或者RoPE旋转位置编码这两种编码对任意长度有更好的泛化性。我建议如果你要部署模型并可能调整分辨率初期就使用相对位置编码或者干脆用2D空间位置编码这样在分辨率变化时能保持坐标关系的连续性。但普通用户直接用官方预训练权重时建议别乱改输入尺寸想用高分辨率就用专门的”动态分辨率”微调策略用小学习率训练一小段时间。4.4 与CNN混合架构的选择ViT把Transformer里最昂贵的全局注意力用在图像每个层上计算量确实吓人。在资源受限或需要实时推理的场景里纯粹的ViT并不总是最优解。混合架构如先用CNN做特征提取再用ViT做深层建模反而是更聪明的选择。具体做法是用ResNet的早期层比如stage1到stage3输出特征图然后把它切割成patch送入几个Transformer编码器最后由分类头输出。这种设计保留CNN的局部先验同时引入Transformer的全局交互训练难度比纯ViT小很多。我在实际项目里用这种混合结构在1万张图片的数据集上就达到跟纯ViT预训练后微调接近的准确率训练时间却少了40%。如果你刚开始接触ViT也可以先用混合架构练手把Transformer部分的原则搞清楚后再挑战纯ViT。5. 常见问题与排查实录最后是实战中的问题排查这部分是我一次次调参、一步步debug积累下来的。如果你在复现或者应用ViT时遇到问题对照这个列表检查一遍能省下大量时间。5.1 训练不收敛先检查数据管道和归一化我见过最多的ViT训练不收敛案例其实模型本身没问题而是数据预处理出了错。ViT对像素值范围敏感必须将输入标准化到标准差为1、均值为0的分布。如果忘记标准化输入范围变成0~255那么注意力点积结果会巨大softmax直接进入饱和区梯度消失。另外如果你的数据是RGB顺序还是BGR也要统一好。我用timm库时甚至发现不同的预训练模型对应不同的mean和std这直接影响微调效果。5.2 显存爆炸的三板斧ViT在中等分辨率下已经够吃显存了如果显存不足最有效的办法不是减小模型而是用梯度累积。梯度累积等价于增大batch size但显存占用不变。混合精度也是必须的能把显存占用砍半。如果你想进一步压就把序列长度降下来把patch size从16改成32序列长度直接变成原来的四分之一。但注意patch变大信息粒度变粗准确率会有所降低所以要从任务出发做权衡。5.3 为什么我的ViT效果还不如CNN排查清单如果你用ViT训完发现效果甚至不如一个ResNet大概率是以下几类问题数据集太小没有使用预训练权重。训练epoch太少ViT需要更长的训练周期。没有参考原论文的训练策略比如学习率太大或太小warmup步数不合理。增强不够ViT强烈依赖数据增强。分类头只训练了最后一层却期望它学到前面的语义特征。我会把我所有的经验浓缩成一句话ViT跟CNN的思维方式完全不同它需要更多的数据、更长的时间、更精细的增强它用这些代价换来了更强的全局建模能力。当你理解了这一点就懂得什么时候该选ViT什么时候该弃ViT了。在视觉项目里模型结构永远是服务于数据和任务的。别为了炫技而硬上ViT也别因为一时的性能差距就完全否定它。如果你愿意接受它的胃口并愿意在工程上做好预处理、预训练、增强和调度它一定会给你带来不少惊喜。这也是我这两年来越来越偏爱它的原因。