视觉Transformer从一出来就带着一股“屠榜”的气势,但真正把它讲明白的文章不多。我最早接触视觉Transformer时,脑子里全是“把图像当句子”这种抽象说法,等自己把一个小模型搭起来跑通之后,才意识到它的核心逻辑其实特别朴素:把图片切成一小块一小块,每个小块当作一个词,然后让模型自己学习小块之间的全局依赖。这篇梳理既是我对自己踩坑过程的复盘,也是给后来者的实操地图。内容涵盖视觉Transformer的架构拆解、最小可复现代码、训练与调参经验、常见问题排查,以及几个最常用的变体对比。适合已经会基本PyTorch、想了解或上手视觉Transformer的同行。
1. 视觉Transformer到底在解决什么问题
1.1 从卷积到自注意力:范式切换的逻辑
在视觉任务里,卷积神经网络统治了很多年,它的强项是局部归纳偏置:用固定大小的卷积核去扫描整张图,天然认为相邻像素的相关性最强,越远的像素影响越弱。这在大多数图像任务上都说得通,所以CNN一直做得很稳。但稳的另一面是“短视”:一个像素要看到整幅图像的信息,必须靠一层层堆叠把感受野慢慢扩大,路径很长,信息在传递过程中也容易丢。
视觉Transformer换了一个思路:不预设局部性,而是让模型自己决定“该看哪里”。它把一张图拆成一个个小块,然后对每个小块做全局注意力计算。所谓自注意力,简单说就是让每一个图像块都去和其他所有图像块做相似度匹配,最后根据匹配结果加权聚合信息。这个过程和图像本身的二维空间排布关系不大,更像是在做一套“谁和谁更相关”的加权投票。
用大白话来类比,CNN像是一个本地关系专家,他跟邻居特别熟,但想了解街区另一头的事得一层层传话;视觉Transformer更像一个全局秘书,它手里的每个人都能直接跟所有人对话,效率高,但前提是它得见过足够多世面,否则很容易把关系搞乱。这个类比能解释为什么视觉Transformer在大数据上越来越强,在小数据上却比CNN更脆弱:归纳偏置少,意味着所有的规则都要靠数据来“喂”出来。
这个范式切换带来的最大优势是长程依赖处理得非常直接:第一层就能看到全局,不用像CNN那样一层层递进。缺点也很明显——它不再自带“邻近优先”的偏置,所有关系全靠数据去学,所以数据量不够时很容易学歪。这也是视觉Transformer在早期被反复质疑“同规模下打不过CNN”的根本原因之一,不是因为注意力机制原理有问题,而是诱导模型学习先验知识的数据量没有给足。
1.2 图像为什么能当作句子处理
NLP里的Transformer处理的是离散的词元序列,每个词元是句子中的一个词。图像在计算机里天然是一个像素矩阵,直接丢给序列模型显然不行。视觉Transformer的做法很直接——把图像分成若干个同样大小的小方块,每个方块展平成向量,当作序列里的一个“词”。这些词的排列顺序由位置编码来表示,顺序信息不能丢。
这里有一个特别好的生活化类比:在一篇文章里,词语的顺序非常重要,“小明打了小红”和“小红打了小明”表达的意思完全不同。图像也一样,一张左边是猫右边是狗的图,和左边是狗右边是猫的图,就是两张完全不同的图。所以视觉Transformer不能只把图像块当成词,还必须给每个块加上表示“我在图像的哪个位置”的位置信息。这也是为什么在Transformer架构里,位置编码不是可选项,而是和图像块向量一样重要的核心组件。
再进一步说,图像其实比句子更好处理,因为它的结构更规则:句子长度有长有短,而图像只要固定输入尺寸,图像块数量就是固定的。这让位置编码的实现可以做得相对统一,也为后面的很多工程优化提供了便利。理解了图像和句子之间“像又不像”的关系,后面看Patch Embedding和位置编码的代码时,就不会觉得它们是什么玄学设计,反而会觉得这些模块的出现几乎是必然的。
需要模型API调用? 免费领10W Token,多模型网关一键接入 Claude、DeepSeek 等主流模型。
2. ViT的核心模块逐层拆解
2.1 Patch Embedding:把图像切成词元
Patch Embedding是整个视觉Transformer的第一步,也是最容易被忽视但很关键的一步。它的输入是一张图像,比如224x224x3,输出是一串二维矩阵,形状是[序列长度, 特征维度]。以默认的16x16图像块为例,224x224的图会被切成196个图像块,每个图像块的大小是16x16x3,展平后就是768个像素值。然后通过一个线性映射把768维压缩到embedding维度,标准配置是768维。
实现上,很多人其实不会真的把图像一块块切出来,再用线性层逐个投影,而是用卷积一步到位:用一个16x16大小、步长也是16的卷积层,直接把整张图卷积成[特征维度, 14, 14]的特征图,再展平转置成[196, 768]。这种方式在工程上速度更快,效果和逐块线性投影完全等价,也是几乎所有开源实现都在用的写法。初次接触的同学看到卷积层突然出现在Transformer前端时容易发懵,记住它只是在做Patch Embedding就行。
除了Patch Embedding之外,序列前面通常还会拼接一个可学习的class token,也就是一个额外的向量,专门用来在最后一层汇总整张图的信息,相当于给模型一个“最终的分类抓手”。同时还需要一个可学习的位置编码,按序列长度加在Patch序列上。这样,整个输入端就把“图是怎么划分的”“每个块在什么位置”“最终输出从哪里取”这三件事安排得明明白白。
2.2 位置编码:给图像块排顺序
位置编码在视觉Transformer里有几种常见形态。最早的工作用的是可学习的一维位置编码:把[196, 768]的序列位置当成一张扁平的查找表,每个位置学习一个768维向量,直接加到Patch Embedding上。这个做法实现最简单,效果也不差,很多开源仓库至今都在用。
后面有研究者提出二维位置编码,也就是分别学习行和列的位置向量,再把两个向量叠加。这么做的好处是模型能更明确地感知图像块之间的行列距离,在做高分辨率迁移时也更符合直觉。实际测试里,二维编码在不少任务上并不会比一维编码有质的飞跃,但可解释性更好,插值时也更稳定,所以在需要调分辨率的下游任务中更受青睐。
这里要特别注意一点:位置编码不是可选的锦上添花,而是决定模型能不能理解图像结构的关键。如果不加位置编码,或者实现有误,模型大概率会退化到“无序词袋”的水平,精度会有肉眼可见的下降。我自己在实验里直接对已经训练好的位置编码做随机置乱测试,分类精度掉得非常厉害,这充分说明Transformer对位置信息的依赖远比很多人想象的要强。Kernel size为16的卷积只是在做空间切分,真正的“空间感”全靠位置编码这个向量来撑着。
2.3 Transformer Encoder:注意力、MLP、LayerNorm
核心编码器的结构如果只挑三样东西讲,那就是多头注意力、MLP前馈网络和层归一化。它们按“注意力-归一化-MLP-归一化”的方式交替堆叠,在视觉Transformer中一般堆12层、24层不等。整个Block的内部逻辑并不复杂,但每一层设计的动机都很值得琢磨。
多头注意力比较像把多个“关系探测器”并联起来。每个头都会把输入分别映射成查询、键、值三组向量,然后计算查询与键的点积得到相似度分数,用softmax归一化成权重,最后加权求和得到输出。多个头各管一摊,有的头关注局部纹理,有的头关注全局轮廓,最后把结果拼起来,相当于让模型用不同视角去看同一张图。点积要做一次缩放,除以键向量维度的平方根,这是为了防止数值过大把softmax推到饱和区,导致梯度消失。
MLP是夹在注意力后面的两个全连接层,中间放一个激活函数。视觉Transformer里默认用GELU,它像ReLU的平滑版本,负数部分不是直接截成0,而是保留很小的梯度,对训练稳定性有帮助。MLP的通道数一般是embedding维度的4倍,比如768维变成3072维再变回768维,这部分占据模型参数和计算量的大头。LayerNorm的作用则是把每一层数据拉回单位尺度,让深层网络的梯度不轻易爆炸或消失。
把这套结构看透彻之后,整个ViT的行为逻辑就通了:注意力负责交换信息,MLP负责变换和记忆,LayerNorm负责稳定训练。每一层做完一次全局信息交换,再经过一次非线性变换,层数越深,信息的组合就越抽象。后面遇到训练不稳定的问题,基本都要回到这三个模块里去排查,而不是盲目调学习率或者堆数据。
3. 手写一个最小可用的ViT
3.1 关键参数与计算量估算
我不打算贴几万行的预训练代码,而是从一个能跑的最小实现出发,把关键参数讲清楚。最常用的基础版配置是:输入224x224,patch_size=16,embed_dim=768,depth=12,num_heads=12,mlp_ratio=4。算下来Patch序列长度是14x14=196,加一个class token,输入序列长度是197。
先估算一下计算量,这样你对模型规模会有直觉。多头注意力的计算量大约正比于序列长度的平方,224分辨率下序列长度只有197,平方还不到4万,相比NLP里动辄几千个token的场景,这个小得可以忽略。所以注意力在这套配置下不是最吃算力的地方,反而是MLP更贵:12层,每层两个大矩阵变换,序列长度197,通道768到3072再到768,这部分占据了绝大部分计算量。这也是为什么把patch从16降到8之后,序列长度从196变成784,注意力部分直接翻了16倍,整体训练时间也随之飙升。
训练层面的超参数设置,在标准配置下其实没有想象中那么玄:预训练常见batch size是4096,学习率取1e-3,用余弦退火衰减,预热5个epoch。如果你的GPU只能跑小batch,学习率也建议等比下调,否则训练很容易起飞。我自己在单卡上训练小模型时,通常把base_lr调到5e-4左右再开始,先跑两三个epoch看曲线形态,再决定要不要微调。
3.2 核心代码结构与训练要点
下面是一个极简但能跑的ViT代码骨架,只保留主干结构,方便直接对照理解:
python复制import torch
import torch.nn as nn
class PatchEmbed(nn.Module):
def __init__(self, img_size=224, patch_size=16, in_chans=3, embed_dim=768):
super().__init__()
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, D, H/P, W/P
x = x.flatten(2).transpose(1, 2) # B, N, D
return x
class Block(nn.Module):
def __init__(self, dim, num_heads=12, mlp_ratio=4.0, dropout=0.0):
super().__init__()
self.norm1 = nn.LayerNorm(dim)
self.attn = nn.MultiheadAttention(dim, num_heads, dropout=dropout, batch_first=True)
self.norm2 = nn.LayerNorm(dim)
hidden = int(dim * mlp_ratio)
self.mlp = nn.Sequential(
nn.Linear(dim, hidden),
nn.GELU(),
nn.Linear(hidden, dim),
)
def forward(self, x):
x = x + self.attn(self.norm1(x), self.norm1(x), self.norm1(x))[0]
x = x + self.mlp(self.norm2(x))
return x
class ViT(nn.Module):
def __init__(self, img_size=224, patch_size=16, in_chans=3, num_classes=1000,
embed_dim=768, depth=12, num_heads=12):
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))
self.blocks = nn.Sequential(*[Block(embed_dim, num_heads) 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, N, D
cls_tokens = self.cls_token.expand(B, -1, -1)
x = torch.cat([cls_tokens, x], dim=1) # B, N+1, D
x = x + self.pos_embed
x = self.blocks(x)
x = self.norm(x)[:, 0] # 只取cls token
return self.head(x)
这个版本用PyTorch自带的MultiheadAttention,省去手写QKV映射的代码,但逻辑上依然对应标准的缩放点积注意力。训练时需要注意几个点:CLS token和位置编码的初始化虽然用零向量也能跑,但实践中用小型正态分布初始化会让收敛更顺畅;dropout和stochastic depth在数据量不大时能明显抑制过拟合。所谓stochastic depth,是每个Block在残差路径上随机丢掉整层的操作,和Dropout作用在神经元上不同,它是作用在层级别的,对Transformer尤其有效。
训练loss建议先用标准的交叉熵。优化器选AdamW而不是Adam,因为解耦权重衰减在Transformer上表现更好。如果要加载预训练权重做迁移学习,强烈建议把分类头换掉后做分层学习率:前面的层级用较小学习率微调,新加的分类头用较大学习率从头训练,这样既保留了预训练特征,又给分类头足够的学习空间。直接全局用一个学习率从头训到底,在小数据集上非常吃亏。
4. 用ViT时最容易踩的坑
4.1 数据量不够怎么办
很多同学发现,同样几千张图,卷积网络能训出不错的验证集精度,ViT却疯狂过拟合,训练集能到99%,验证集只有六十几。这不是bug,而是模型特性。Transformer的全局注意力参数自由度更高,需要大量数据才能把这些自由度约束到有用的方向上。图像块之间不再是局部物理邻近,而是完全靠数据学出来的关系权重,样本少时很容易记住训练集里的噪声模式。
如果只有几千张图,我通常先做三件事。第一,换小patch或者加卷积茎干,把patch从16降到8,相当于把更多卷积式的局部归纳偏置引入模型,同时序列变长后模型会有更多位置感知的空间。第二,在分类任务的标准数据增强上加大力度,混合增强和随机擦除至少开一个,再配合随机增强,直接把有效样本量放大一个量级。第三,想尽办法拿预训练权重,哪怕拿到的预训练任务和你的任务不完全一致,用它的前几层做初始化也比随机初始化稳得多。
如果以上都试完还是不行,就不要硬上纯ViT了,换成带层次结构的变体,或者直接上带卷积茎干的混合半变体,几乎都能缓解过拟合。我的一个经验是:视觉Transformer在小数据集上不是不能用,而是必须搭配足够强的约束,指望靠训练技巧把随机初始化的ViT在几千张图上训好,性价比非常低。
4.2 训练不稳定与学习率策略
ViT的训练曲线经常出现一种让人头疼的现象:loss掉了几个epoch之后突然弹一下,然后整个训练开始震荡。我会第一时间查三个地方:学习率是不是太大、预热是不是太短、梯度裁剪有没有开。Transformer的注意力层对学习率比卷积更敏感,大学习率会让QKV映射在早期游走得太远,等预热结束进入主训练阶段时,模型已经站在了不稳定的山坡上。
稳妥的做法是:AdamW优化器配1.0的梯度裁剪,学习率走余弦退火,并且把warmup拉长到总步数的5%。比如总训练100个epoch,预热5个epoch比较合理,太短等于没预热。Batch size调大的时候,学习率可以按线性缩放规则近似放大,但不要放大超过原值的3到4倍。我用单卡batch 64训练小ViT时,base_lr一般取3e-4到5e-4,很少直接套用大batch配置里的1e-3,实测直接套大学习率几乎必炸。
还有一个容易被忽略的是stochastic depth的配置。它相当于给每个Block增加一个随机丢层概率,像Dropout但作用在层级别。对ViT来说,这个正则化手段远比常见的embedding dropout重要,因为在深层Transformer里,某一层过拟合对整体影响特别大。配置一般在0.1到0.3之间,数据越少丢层概率越高。我的习惯是先固定学习率策略,再调这个丢层概率,把它当作最强力的过拟合刹车来用。
4.3 位置编码与分辨率迁移问题
训练好一个224分辨率的ViT后,想把它用在384或者更大的输入上,直接把位置编码的长度改成新序列长度是行不通的,因为位置编码是可学习参数,新长度对应的位置从来没被训练过。常见做法是把已有的位置编码做插值:把14x14的位置网格双线性插值到24x24的网格,再展平使用。实操里,二维插值比一维插值效果更好,因为它保留了图像块的行列空间关系,一维扁平插值会破坏这种结构。
但插值不是无损的,高频信息在插值过程中会被平滑掉。我的经验是,插值之后最好做一小段微调训练,让模型重新适应新的位置关系,直接拿去推理的效果通常不如训练时就用目标分辨率。还有一个经常被问的问题:分辨率提高后,patch_size要不要跟着变?常规做法是不变,patch_size保持16,分辨率从224提到384,序列长度从196变成576,计算量会显著上升。如果算力紧张,可以把patch_size同时提高到24或者32,但这会增加信息丢失的风险,属于取舍问题。
最后提一个很少出现在文档里的细节:位置编码的参数初始化和尺度也会影响迁移效果。某些开源实现里位置编码用截断正态分布,而有的直接用零初始化,差别在从头训练时没那么大,但在微调大分辨率模型时,尺度不匹配会导致插值后的特征出现偏移。如果遇到“预训练权重在224表现很好,插值到384后验证曲线完全崩掉”的情况,可以先试试只冻结位置编码、微调其他部分,或者把位置编码学习率设得特别小,很多情况下能救回来。
5. 从ViT到它的变体们
5.1 DeiT:蒸馏与数据效率
ViT在大数据上很强,但小数据集上一直被CNN压制,一个很重要的改进方案是DeiT。它的核心思路是:在训练时额外引入一个教师模型,通常是卷积网络,让ViT不仅学习真实标签,还学习教师的输出分布。具体做法是在class token旁边再插一个distillation token,训练时用蒸馏损失约束这个token的输出,让学生的行为尽量向教师的预测靠拢。这个token和class token在推理时既可以合并使用,也可以单独取class token。
这个思路之所以有效,是因为教师模型已经把所有视觉经验编码在它的输出概率里,学生不用再从头摸索图像块之间的复杂依赖关系。DeiT用一半的训练数据就能追平原始ViT在更大数据集上训练的效果,足见知识蒸馏对Transformer的友好程度。如果你拿到一个小任务又想坚持用Transformer,找一个强一点的CNN或者更大的模型当老师,基本等于白捡一大截精度。
实操里有几个细节要注意:蒸馏温度一般设置在3到5之间,太高会把概率分布压得太平,太低又学不到教师的暗知识。distillation token和class token的初始化方式、学习率设置也要稍微区分,两个token本质上承担了不同任务。如果你用的框架只实现了标准ViT,临时加一个蒸馏token并不难,代码层面其实就是多个拼接向量和一组额外的损失计算。
5.2 Swin Transformer:层次化与窗口注意力
Swin可能是视觉Transformer家族里最偏向工程落地的一个。它的核心改动有两个:一是借鉴CNN的分层结构,让特征图尺寸逐渐从整图缩小到1/2、1/4、1/8,形成多尺度特征金字塔,这对检测、分割这类需要在不同尺度上做预测的任务特别关键;二是把全局注意力改成窗口注意力,只在固定大小的窗口内做自注意力,窗口之间再通过一次移位操作来传递信息。
这个设计的聪明之处在于计算复杂度:全局注意力是序列长度的平方,而窗口注意力的复杂度随图像尺寸线性增长。同样是224输入、patch_size为4的设定,Swin的计算量比标准ViT低得多,因此可以在更大的图像上运行,也可以做更深的网络。窗口移位则解决了纯局部注意力信息不互通的问题,相邻两次窗口划分错开半个窗口,每个图像块都有机会和远处的图像块产生关联,效果上接近全局注意力但代价小很多。
实际用下来,Swin在做下游任务时明显比标准ViT更顺手。它天然有多尺度特征,目标检测和图像分割框架可以直接把它当backbone接入,几乎不需要额外修改。如果项目里用的是检测或分割模型,想换成Transformer架构,我建议第一个看Swin而不是标准ViT,能省下非常多适配时间。
5.3 其他方向:Patch的尺度、移位、混合架构
除了DeiT和Swin,整个领域还衍生了非常多有意思的方向,我自己关注比较多的有三类。第一类是Patch尺度的探索。之前说过patch_size决定了序列长度和计算量,有不少工作专门研究怎么在不同层或者不同空间位置使用不同大小的patch,让模型可以像人一样先粗后细地看图。第二类是把局部注意力和全局注意力交替排列,在保持效率的同时尽量不丢全局信息,这类设计在分割和目标检测任务上也很常见。第三类是混合架构,把卷积和Transformer放在同一个模型里,前面几层用卷积做降采样和特征提取,后面几层用Transformer做全局建模,这种做法在图像分类上的效果直逼纯Transformer,但训练更稳定,数据效率也更高。
这些方向本质上都在解决同一个问题:如何在享受全局建模能力的同时,尽可能压低计算量、提升数据利用效率。后续如果你的业务场景对计算设备有严格限制,建议从这些混合或层次化模型入手,别一上来就硬扛标准ViT的算力需求。
最后说一点我在实际项目里的体会。视觉Transformer并不是什么神秘魔法,它的出现把计算机视觉的很多任务统一到了“序列建模”的框架里,但这也意味着你不能再像以前训CNN那样无脑堆经验。最重要的三个准备工作——数据量评估、预训练权重获取、训练超参数设计,每一步都值得单独花时间。如果只记一句话,那就是:先保证训练能稳定收敛,再谈效果和涨点。这个小经验,帮我省下了无数个深夜排查训练曲线的周末。
