简介面向PyTorch用户的Vision TransformerViT实现包由原始JAX/Flax预训练权重转换而来适合需要在PyTorch生态中做图像分类、特征提取或下游任务微调的研究者与开发者。包内共35个文件以Python源码为主22个py覆盖模型定义、训练、评估、数据加载、配置与工具脚本辅以Markdown文档、notebook示例、yml配置和license说明目录结构清晰便于按需查找。整个压缩包仅173KB纯代码形式可快速阅读与改动。目前已有7854人学习浏览。资源提供从JAX/Flax权重转换的预训练ViT模型支持ImageNet2012等数据集并附带微调与评估脚本以及可直接运行的示例notebook能帮助用户快速复现论文结果或迁移至自有视觉任务适用于有PyTorch基础、希望使用ViT的开发者。1. 先弄清楚再动手ViT 是什么这份代码能帮你省掉什么先讲一个反直觉的观察Vision Transformer 在没有卷积核的情况下靠“切块 → 嵌入 → Transformer 编码”这套流程在 ImageNet 上追平了同体量 CNN但真正让它在工程里可用的不是模型定义那几十行代码而是预训练权重的完整度和加载逻辑的健壮性。vision-transformer-pytorch 是某开源组维护的 PyTorch 实现把 ViT 拆成标准 nn.Module 模块并把预训练模型直接接进加载流程省掉你手动找权重、对齐 state_dict、配归一化参数这些脏活。适合两类人一是着急在自己数据集上做微调的工程师二是刚读完 ViT 论文、想逐行对着代码验证理论的学生。下面按“结构 → 推理 → 踩坑 → 微调”逐层拆开说。2. ViT 拆解补丁化、注意力头与序列长度如何决定模型容量上手这份代码之前建议先把一层层的张量形状捋清楚。ViT 的输入是一个标准图片张量但模型内部处理的是 token 序列——这个思维转换决定了你后面能不能改得动别人的代码。2.1 补丁化一张图如何变成 token 序列我第一次手写 ViT 的时候卡了很久图像张量是 [B, 3, H, W]Transformer 期望的是 [B, seq_len, dim]中间这条桥怎么搭答案是补丁化Patch Embedding。把 224×224 的图按 16×16 的窗口切成互不重叠的小块得到 14×14196 个 patch。每个 patch 展平后是 3×16×16768 个数值恰好映射到嵌入维度 768。这就是为什么 ViT-Base 的 dim 是 768 而不是拍脑袋定的。实际工程实现里补丁嵌入几乎都是用卷积一步完成而不是真正去“切块”import torch import torch.nn as nn patch_embed nn.Conv2d( in_channels3, out_channels768, # 对应 dim即 token 嵌入维度 kernel_size16, # patch 边长 stride16 # 步长等于 patch 尺寸窗口不重叠 ) x torch.randn(4, 3, 224, 224) # 4 张 RGB 图每张 224×224 x patch_embed(x) # 输出 [4, 768, 14, 14] x x.flatten(2) # 拉平空间维度 - [4, 768, 196] x x.transpose(1, 2) # 转成 Transformer 习惯的 [4, 196, 768]卷积核在每个 16×16 窗口内做加权求和等价于把每个 patch 展平后乘一个 [768, 768] 的线性矩阵所以用卷积代替手写循环结果完全一样但底层算子优化更好。这里有个参数细节要盯紧image_size 必须能被 patch_size 整除。你把输入改成 128×128 而 patch_size 保持 16那 token 数变成 64 个如果改成分辨率 130flatten 出来就是一个非整数的空间尺寸位置嵌入直接对不上。2.2 位置嵌入和 CLS token197 这个数字怎么来的Transformer 本身没有位置概念QKV 注意力在计算时对序列顺序不敏感所以必须显式把位置信息注入输入。ViT 的做法是初始化一个可学习的位置嵌入矩阵而不是用固定的正弦编码形状是 [1, 197, 768]其中 197 196 个 patch token 1 个 CLS token。对应操作如下# 可学习的 CLS token整个序列只加一个 cls_token nn.Parameter(torch.randn(1, 1, 768)) # 位置嵌入197 个位置每个位置 768 维 pos_embed nn.Parameter(torch.randn(1, 197, 768)) # 把 CLS token 拼到 patch token 序列头部 x torch.cat([cls_token.expand(B, -1, -1), x], dim1) # [B, 197, 768] # 逐位置相加 x x pos_embed这里容易忽略的细节是pos_embed 的形状必须和实际 token 数完全一致。改 patch_size 或输入分辨率位置嵌入尺寸就会不一致加载预训练权重时报 shape mismatch 是必然的。这也是微调低分辨率图像比如 128×128时最常见的报错来源。还有个隐藏点这个仓库里 CLS token 只在序列头部放一个而不是每个 patch 都加理解错这个结构的人改代码时会越改越晕。训练完成后分类只用 CLS token 对应的输出向量。这和 NLP 里 BERT 的做法一脉相承Encoder 在训练过程中会把全局类别信息逐步汇聚到这个 token 上。为什么不用平均池化因为平均池化会把不同 patch 的局部信息无差别混在一起而 CLS token 相当于让模型自己学出来一个“全局语义汇聚点”。2.3 Transformer 块前归一化、多头注意力与 MLP 的职责分配每个编码器块做两件事先通过多头自注意力做 token 间的信息交换再做 MLP 做逐 token 的非线性变换。ViT 基本都采用前归一化Pre-LN结构也就是 LayerNorm 放在子层之前。这样做梯度流更稳也允许你用更大的学习率class Block(nn.Module): def __init__(self, dim768, heads12, mlp_dim3072, dropout0.): super().__init__() self.norm1 nn.LayerNorm(dim) self.attn nn.MultiheadAttention(dim, heads, dropoutdropout) self.norm2 nn.LayerNorm(dim) self.mlp nn.Sequential( nn.Linear(dim, mlp_dim), nn.GELU(), nn.Dropout(dropout), nn.Linear(mlp_dim, dim), nn.Dropout(dropout), ) def forward(self, x): # 残差连接 前归一化先 norm 再进子层 x x self.attn(self.norm1(x), self.norm1(x), self.norm1(x))[0] x x self.mlp(self.norm2(x)) return x注意力机制里Q、K、V 三个矩阵都由输入 x 经过不同的线性投影得到注意力权重就是 Q 和 K 的点积加 softmax。多头的作用是让不同的头关注不同的模式——有的头看全局轮廓有的头看局部纹理最后拼接起来信息维度更丰富。以 ViT-Base 为例dim768heads12每个头处理 768/1264 维子空间。所有头的输出拼回去还是 768 维经过一层输出投影再进残差。网络主体由 depth 个这样的 Block 堆叠ViT-Base 是 12 层。在这份代码里修改 depth 和 dim 就能得到更小或更大的变体比如想要 ViT-Small 就把 depth 设成 8、dim 设成 384但要注意没有官方预训练权重你需要自己从头训练。从参数量来看各组件差异很大组件参数量约说明补丁嵌入层59 万3×768×16×16 卷积权重位置嵌入 CLS token15 万197×768 76812 个编码器块约 8500 万每个块约 710 万含注意力与 MLP分类头 最终 LayerNorm约 77 万768×1000 768×2ViT-Base 总参数量大约 8600 万比 ResNet-50 大了一倍还多。模型容量大需要更多的数据才能发挥威力。几千张图从头训练 ViT 大概率不如同体量 CNN 微调用预训练权重做微调才是正路。3. 加载预训练权重跑推理预处理、权重校验与分类头替换3.1 实例化模型pretrained 参数背后发生了什么Vision Transformer 的推理入口很简洁但参数含义值得逐个对一遍import torch from vit_pytorch.vit import ViT model ViT( image_size224, # 输入分辨率必须能被 patch_size 整除 patch_size16, # 切块尺寸 num_classes1000, # 预训练权重对应的类别数 dim768, # token 嵌入维度 depth12, # Transformer 编码器块数 heads12, # 多头注意力头数 mlp_dim3072, # FFN 中间层维度 pretrainedTrue, # 触发预训练权重加载 ) model.eval()pretrainedTrue 一般会触发权重下载。仓库按预设 URL 去拉取权重文件存到本地缓存目录下次启动直接读缓存。第一次下载比较慢而且受网络环境影响大。更稳妥的做法是手动下载 .pth 文件放进预期路径或者直接用 load_state_dict 加载本地文件。这里最容易被坑的参数是 num_classes。预训练权重在 ImageNet-1k 上训练输出是 1000 类。必须先让模型结构对齐权重结构num_classes1000加载权重然后再改分类头。如果你一上来就设成自己数据集的类别数模型和权重结构不一致加载必然报 size mismatch。3.2 预处理与推理ImageNet 归一化参数为什么不能省用预训练权重做推理输入图像的分布必须还原训练时的样子。ImageNet 数据集的统计均值是 [0.485, 0.456, 0.406]标准差是 [0.229, 0.224, 0.225]。用错归一化参数模型输出基本等于随机——这是新手最常见的“为什么效果这么差”的原因。from PIL import Image from torchvision import transforms transform transforms.Compose([ transforms.Resize(256), transforms.CenterCrop(224), # 和预训练输入对齐 transforms.ToTensor(), transforms.Normalize( mean[0.485, 0.456, 0.406], std[0.229, 0.224, 0.225] ), ]) def predict(image_path, model): img Image.open(image_path).convert(RGB) # 强制 RGB过滤灰度/透明通道 x transform(img).unsqueeze(0) # [1, 3, 224, 224] model.eval() with torch.no_grad(): logits model(x) prob torch.softmax(logits, dim-1) return prob[0] prob predict(demo.jpg, model) topk torch.topk(prob, k5)为什么先 Resize 到 256 再 CenterCrop 到 224而不是直接 Resize 到 224因为训练时用的是随机裁剪图像先缩放到 256 保证有足够边缘信息CenterCrop 取中间区域是推理时固定化的近似。如果你直接 Resize 到 224高宽比被拉伸变形模型分类置信度会明显下降。这个细节在 ImageNet 数据集上能差出两三个点的 top-1。3.3 替换分类头迁移到下游任务的第一步在自己数据集上做分类类别数大概率不是 1000。常规做法是把最后的分类头换掉num_classes_self 17 model.fc torch.nn.Linear(768, num_classes_self) # 有的仓库里这个属性叫 model.head替换前先打印 model 确认 state torch.load(vit_base_patch16_224.pth, map_locationcpu) model.load_state_dict(state, strictFalse)strictFalse 的意思是允许 state_dict 里缺 key 或有多余 key。但这里的风险在于——它只是掩盖错误不是修复错误。一个我保留至今的排错习惯是打印缺失的 keymissing_keys, unexpected_keys model.load_state_dict(state, strictFalse) print(缺失的 key:, missing_keys) print(多余的 key:, unexpected_keys)如果缺失列表里出现 pos_embed、blocks.0.attn.in_proj_weight 这类主干层说明你的模型结构和预训练权重根本对不上strictFalse 只是把错误藏起来了。如果缺失的只有分类头的两个 keyweight 和 bias才是正常替换。加载完成后跑一次样例图如果 top-1 接近随机先怀疑预处理别怀疑权重。4. 避坑手册ViT 加载与微调的五个高频故障4.1 权重下载失败和 state_dict 加载报错现象一自动下载权重时网络超时或中断。程序抛出 HTTP 错误或者下载完发现文件只有几 KB加载时报“unexpected end of file”之类的损坏提示。原因自动下载走的是仓库预设 URL过程没有断点续传也没有文件完整性校验。网络波动时很容易拿到半截文件。解决手动下载权重到本地对比官方标注的 MD5 校验值确认完整。然后把文件放进代码期望的缓存路径或者直接用 load_state_dict 从本地路径加载。加载完顺手打印一个 block 的权重均值和方差确认不是随机初始化残留。现象二加载报错 size mismatch for fc.weight。错误内容通常是“copying a param of shape torch.Size([1000, 768]) to …”也就是分类头维度对不上。原因分类头已经被你换成自定义维度的线性层或者 num_classes 参数一开始就设成了自己的类别数。两种情况的本质都一样模型最后线性层的输出维度与预训练权重不一致。解决先保持 num_classes 和预训练权重一致加载加载成功后再替换分类头。替换后如果还要加载同一个权重文件用 strictFalse 且确认缺失 key 只有分类头相关的那两个。4.2 预处理不一致导致推理结果像随机现象三同一张图官方示例脚本能预测正确换了自己的预处理代码后 top-1 明显不对甚至输出概率分布平坦。原因尺寸处理和归一化参数不一致。常见错误有三种直接把图 Resize 到 224×224 导致高宽比变形换了自己数据集的 mean/std把图片先乘 2 再减 1以为是 [-1,1] 归一化。预训练权重训练时看到的是固定输入分布推理不做前置对齐效果跟随机初始化几乎一样。解决统一走 Resize(256) → CenterCrop(224) → ToTensor → Normalize(ImageNet 参数) 这条流程。凡是出现“官方示例能过、自己数据不 work”先逐行核对前四步是不是完全一致再考虑模型问题。现象四批处理某些图片时崩溃报错和通道数有关。提示你的输入是 [4, H, W] 或 [1, H, W]模型期望 [3, H, W]。原因数据里有带 alpha 通道的 PNG或者灰度图、损坏图混进来了。PIL 打开后默认保留原始通道不做转换就进 ToTensor通道数自然对不上。解决统一用 .convert(RGB) 强制转三通道。批处理前先跑一遍数据完整性检查用 Image.open img.verify() 把损坏文件提前筛出来别等循环跑一半才崩。4.3 显存 OOM 和验证指标诡异现象五batch_size 调到 32 直接 CUDA out of memory降到 8 也经常崩。原因Transformer 的自注意力要维护 [B, heads, seq_len, seq_len] 的注意力矩阵。以 ViT-Base 为例12 个头 × 12 层内存占用随序列长度平方增长。如果把 patch_size 从 16 改成 8token 数变成原来的 4 倍注意力矩阵就是 16 倍显存直接爆炸。解决优先用梯度累积来模拟大 batch而不是直接调大 batch_size。训练时开 AMP自动混合精度能省近一半显存。还不行就换回 patch_size16 或调小输入分辨率——但改分辨率连带要处理 pos_embed 插值比较麻烦。动手前先 nvidia-smi 看一眼显存占用之前我遇到过工作站的显存被其他进程占了三分之一查了半天才发现是环境问题。验证集指标接近随机、loss 不降先别急着怀疑代码逻辑。按顺序排查打印模型 logits 数值范围是否正常用预训练权重跑一张 ImageNet 样例验证管线最后才检查自己的数据加载和标签对齐。大部分“玄学”问题都出在预处理和标签错位这两块。5. 微调调参框架冻结策略、优化器与数据增强的搭配5.1 冻结策略从特征提取到全量微调的梯度释放预训练权重已经是在 ImageNet 上收敛过的解不需要一上来就调所有参数。数据量小的时候全量微调容易过拟合冻结主干只训练分类头是最稳妥的起点。# 第一步全部冻结 for param in model.parameters(): param.requires_grad False # 第二步解冻最后 4 个编码器块 for i in range(model.depth - 4, model.depth): for param in model.blocks[i].parameters(): param.requires_grad True # 第三步分类头保持可训练 for param in model.fc.parameters(): param.requires_grad True代码里的 blocks 和 fc 命名在不同仓库里略有差异有的是 blocks有的是 encoder.layer。拿到代码先打印 model 看一遍模块名别凭记忆写。深度为 12 的 ViT-Base从第 8 个块开始解冻相当于只让最后三分之一参与更新前面保持预训练特征不变。冻结层数的选择有几个经验值几百到几千张图只解冻最后 4 个块加分类头数据量上万可以全量微调但学习率要调低一个量级如果是医学影像这类和 ImageNet 分布差异很大的数据冻结策略收益有限直接全量微调通常更好。核心原则是数据越少更新参数越保守。5.2 优化器与学习率ViT 对学习率的敏感度高于 CNNViT 微调最常见的优化器选择是 AdamW它对 Transformer 结构更稳定权重衰减和大模型结合得也好。学习率起步区间建议 3e-5 到 1e-4不要一上来就 1e-3。optimizer torch.optim.AdamW( filter(lambda p: p.requires_grad, model.parameters()), lr5e-5, weight_decay0.05, ) scheduler torch.optim.lr_scheduler.CosineAnnealingLR( optimizer, T_max30, # 总训练轮数按任务调整 eta_min1e-6, )从零训练 ViT 时学习率可以往 1e-3 走但微调时权重已经在最优解附近学习率需要降两个数量级在局部极小点附近小步调整。判断标准很简单第一个 epoch loss 不降反升几乎可以断定学习率太大。ViT 对学习率的敏感度比 ResNet 高一个等级翻车概率也大不少。我有一次把 lr 设到 3e-4 微调loss 直接起飞降到 5e-5 才恢复正常。从那以后每次换数据集我都会先用 5e-5 跑 3 个 epoch 看趋势再决定加减。调度器别用固定学习率。余弦退火是 ViT 微调的默认选项它在前半程保持较高梯度、后半程平滑收敛不容易踩到陡峭损失面。5.3 数据增强小样本下成本最低的止损手段ViT 相比 CNN 在小样本场景下对过拟合更敏感而数据增强是投入产出比最高的修正手段。常用方案分两块from torchvision import transforms train_transform transforms.Compose([ transforms.RandomResizedCrop(224, scale(0.05, 1.0)), transforms.RandomHorizontalFlip(), transforms.ColorJitter( brightness0.4, contrast0.4, saturation0.4 ), transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]), ])RandomResizedCrop 的 scale 下限设到 0.05意味着最小只保留画面的 5%强制模型学习关注局部纹理。这个参数是我在一次 2000 张图的数据集上调出来的比 scale 默认 0.5 高了好几个点。ColorJitter 的幅度不要贪大亮度对比度饱和度的幅度在 0.3 到 0.4 之间比较安全超过 0.5 容易把颜色当成噪声丢掉。再往上一层可以考虑 MixUp 和 CutMix。ViT 的注意力机制天然适应这两种增强因为模型能学会区分混合区域的语义权重而不是简单地把标签当成噪声。如果数据不平衡加权采样器比在 loss 里加权重更直接有效。验证集的预处理必须和第 3 章的标准流程保持一致不能在验证集上随机裁剪。这是微调里最常见的隐蔽错误训练集做了 RandomResizedCrop验证集也照抄了同一个 transform导致验证指标虚高部署时直接打回原形。6. 进阶技巧用注意力图给模型做体检6.1 一行 hook 提取注意力分布预训练权重对不对、预处理有没有问题不需要等训练跑完看注意力图心里就有数。ViT 的注意力矩阵本身就可解释attn_map [] def hook_fn(module, input, output): # attn 模块返回 (attn_output, attn_weights)取第二个即可 attn_map.append(output[1].detach().cpu()) # 取最后一个编码器块的注意力层 model.blocks[-1].attn.register_forward_hook(hook_fn) with torch.no_grad(): _ model(x) # 形状 [B, heads, 197, 197]取 CLS token 对所有 patch token 的关注权重 cls_attn attn_map[-1][0, 0, 0, 1:].reshape(14, 14)把 cls_attn 上采样回 224×224 叠加到原图上就能看到模型在关注哪里。这里的 14×14 恰好对应 patch_size16 时的网格数量224/1614。如果你想看不同层的注意力聚合就把多层的结果取平均想看不同头的行为就取 [0, 4, 0, 1:] 改一下索引。6.2 什么样的注意力是健康的一张正常图上注意力应该有明显的局部聚集目标物体在画面中央区域注意力权重集中在目标边界和纹理处背景区域的权重趋近于零。如果注意力图像撒胡椒面一样无差别铺满整个画面先怀疑预处理归一化出错再怀疑权重加载错位。如果注意力集中在四个角落大概率是 CenterCrop 没做模型看到了自己的 padding。我之前有一次微调 128×128 输入加载权重时报错一偷懒把 pos_embed 直接裁掉了一段继续训练结果前几个 epoch 指标正常后面越训越乱。后来用注意力图检查才发现位置信息全乱了。从那以后我每次拿到新的预训练权重都强制走一遍单图预测加注意力图可视化的流程确认输入预处理、权重对齐、类别索引三层都通了才开始微调。这个检查用不到两分钟但能省掉后面几天白跑的训练。希望帮到你。本文还有配套的精品资源点击获取