简介这是一份基于 PyTorch 的 Vision TransformerViT实现面向深度学习研究者与工程师提供从原始 JAX/Flax 权重转换而来的预训练模型可直接用于图像分类、特征提取以及下游任务微调。压缩包约 173KB共 35 个文件其中以 22 个 Python 源码为主涵盖模型定义、训练、评估、数据加载与配置管理另有 README、Markdown 说明、requirements 环境文件、YAML 配置及 Notebook 示例等目录结构清晰便于二次开发。资源描述了与原始模型相当的结果支持 ImageNet2012 等数据集并附带微调与评估脚本适合有一定 PyTorch 基础、希望快速复现 ViT 论文或将其应用于自身视觉任务的开发者。已有 7800 余人浏览学习是入门 Vision Transformer 并获取可用预训练权重的实用参考。 我做视觉这两年最常用的框架就是PyTorch而ViT相关的项目里vision-transformer-pytorch这个库我几乎是反复在用。很多朋友第一次接触Vision TransformerViT时被论文里一堆概念唬住总觉得这是个很复杂的模型。但如果你上手跑一遍这个库会发现ViT的思路其实非常直接把图像切成小块当成一串“视觉单词”送给Transformer去处理。而且这个项目坐标很明确——Pytorch 预训练模型正好是现在视觉任务落地最常用的一套组合。这个项目解决的核心问题就是让大家不用重复造轮子。它提供了完整的ViT模型实现结构清晰、参数可调并且可以配合预训练权重直接使用。对我个人来说它最大的价值在于既能帮新手理解ViT内部到底发生了什么也能让有经验的工程师在几行代码内完成模型搭建和迁移学习。无论你是打算做图像分类还是把ViT当作骨干网络接进检测分割框架这篇文章都适用。1. ViT模型为什么值得关注从CNN到Transformer的范式转移1.1 图像能否直接当序列处理在ViT出现之前视觉模型几乎被卷积神经网络CNN统治。CNN的核心假设是局部性和平移不变性相邻像素关系更密切同一个卷积核在整张图上滑动。这个假设在ImageNet这类中大规模数据集上非常好用因为卷积天然带有先验不需要太多数据就能学会。但Transformer的想法完全不同。它最初在NLP里证明了一件事只要数据足够多你不给模型任何结构先验让注意力机制自己去找全局关系效果反而可能更好。于是就有了一个自然的问题图像能不能也当成一个token序列来处理ViT回答这个问题的方式很直接——把图像切成固定大小的patch每个patch线性映射成一个向量再按顺序拼起来当作一个“句子”来处理。我最早看到这个思路的时候也挺惊讶整个ViT居然没有一个卷积层全靠自注意力在建模。也正是这种极简让它在超大数据集上表现惊人。论文里用JFT-300M这种量级的数据集做完预训练之后ViT在ImageNet上的精度能超过同等规模的ResNet和EfficientNet。所以它不是换了个结构而是换了一种视觉建模的范式。1.2 vision-transformer-pytorch解决的三个实际问题用这个库一年多我感受到的突出价值有三点。第一实现与论文高度对齐可读性强。它不是把ViT封装成黑盒而是把patch embedding、transformer encoder、分类头拆成模块改哪里都一目了然。我在对比不同层数、不同head数量对精度影响的时候基本就是改几个参数的事。第二原生PyTorch生态无缝衔接。库本身不依赖timm或者更上层的框架直接就能和你自己的训练管线集成。我经常需要把ViT输出的特征接给检测头或者分割头这个库的数据流非常透明改起来很顺手。第三社区认可度高坑少。这个项目在GitHub上star量很大用的人多意味着踩坑经验多。比如位置编码的维度问题、patch size选择问题网上一搜就有很多讨论遇到bug不至于孤立无援。对做工程和做研究的人来说这种成熟度很重要。2. 核心架构拆解ViT到底在做什么2.1 Patch Embedding图像是如何变成Token序列的ViT最核心的一个操作就是Patch Embedding。以最常见的ViT-B/16为例输入是224x224的RGB图像patch_size设为16那么图像会被划分成(224/16)(224/16)1414196个patch。每个patch大小为16x16x3把它展平成长度为768的向量再经过一个线性投影层映射到768维的embedding空间。很多人第一次看会疑惑为什么不直接展平再用全连接其实这里的线性投影本质就是一个1x1卷积或者一个reshape加Linear它的作用是让每个patch的原始像素映射到更适合Transformer处理的语义空间。实际操作中很多实现直接用nn.Conv2d(in_channels3, out_channelsdim, kernel_sizepatch_size, stridepatch_size)来一步完成切patch和投影效率更高。这也是为什么你会看到有些代码里Patch Embedding层长得像卷积层但它做的事情其实就是“切块线性变换”。这里我想强调一个点patch size的选择直接影响序列长度。patch越小序列越长计算量越大但细节保留越多。ViT-B/32用32的patch序列长度只有491个token速度快很多但精度略降。实际工程里如果显存有限又不想掉太多精度可以考虑用大patch或者保持patch不变减少层数。2.2 位置编码、CLS Token与Transformer Encoderpatch被映射成token之后接下来的问题很关键Transformer本身是顺序无关的它不知道哪个token在图像的哪个位置。ViT的做法是加一个1D可学习的位置编码向量直接加到所有token的embedding上。这里没有用NLP里常见的2D位置编码因为论文实验发现1D可学习编码对效果影响不大但实现更简单。ViT还在序列最前面插入了一个特殊的CLS token它的作用和BERT里的CLS一样用于汇聚全局信息。在Transformer编码若干层之后模型拿出CLS token对应的输出向量接一个分类头完成最终分类。我实际中还见过一些改造方案比如直接对所有token做全局平均池化再分类效果有时也不差但标准ViT用的是CLS token方案。接下来是Transformer Encoder。以ViT-Base为例包含12层Encoder每层由多头自注意力12个head、MLPhidden size从768扩展到3072再降回来、LayerNorm和残差连接组成。值得注意的是ViT用的是Pre-LayerNorm结构也就是每个子层注意力或MLP之前先做归一化。这个细节影响稳定性训练大模型时尤其明显。我自己的经验是这种设计配合较大的学习率也能保持稳定微调时不容易崩。2.3 模型规格怎么选Base、Large与HugeViT官方发布了几个规格Base86M参数、Large307M参数、Huge632M参数。视觉任务里用最多的就是Base它和ResNet50规模差不多但效果更好。如果资源充足、任务复杂Large往往能带来明显提升Huge则适合在超大数据集上从头训练一般做迁移学习的用不起。我选型时通常会先想清楚数据集规模。数据集只有几千张图直接用Base甚至Small版本配合在ImageNet上预训练的权重效果往往比从零训练要稳得多。所以这里就引出了下一部分的重点预训练模型到底怎么选、怎么用。3. 预训练模型的选择与微调实战思路3.1 三种获取预训练权重的方式对比标题里特别提到了“带有预训练模型”这一点其实是很多人最关心的。vision-transformer-pytorch库本身侧重于提供模型结构而预训练权重通常可以从下面三个渠道获取获取渠道是否携带官方预训练权重适用场景vit-pytorch库否只提供模型结构学习结构、自定义改造timm是ImageNet预训练权重日常分类、工程落地HuggingFace transformers是Google官方权重研究复现、需要官方预处理我个人的建议是追求简单就直接用timm一行代码搞定下载和加载追求跟原论文对齐就去HuggingFace。用timm加载预训练权重的方法非常直接import timm model timm.create_model(vit_base_patch16_224, pretrainedTrue, num_classes1000) model.eval()如果你要做的分类任务类别数不是1000可以直接在create_model时指定num_classestimm会自动把最后的分类头替换成对应数量微调时非常省事。另外timm还提供了丰富的数据增强策略、EMA等训练工具对训练精度的提升很友好。如果你更希望复现原论文的预处理流程可以用HuggingFace的transformersfrom transformers import ViTForImageClassification, ViTImageProcessor model ViTForImageClassification.from_pretrained(google/vit-base-patch16-224) processor ViTImageProcessor.from_pretrained(google/vit-base-patch16-224)这里的processor封装了标准化和尺寸调整逻辑拿过来就能用不用自己纠结normalize参数。不过要注意transformers加载这种大权重时会从HuggingFace服务器下载虽然大部分时候没问题但偶尔会卡住这时候可以设置环境变量HF_ENDPOINThttps://hf-mirror.com用国内镜像加速下载。注意预训练权重的下载只是第一步真正决定模型效果的是后续微调策略。3.2 微调策略冻结层、学习率与数据规模拿到预训练模型后第一个选择就是冻结还是不冻结。如果数据集比较小、只有几千张我建议冻结前11层Encoder只微调最后一层和分类头。实操中Encoder底层学习到的都是一些基础纹理、边缘特征这些对于任何视觉任务都是通用的不需要重新学。冻结之后可以大幅减少显存占用和训练时间加快收敛。如果数据集中等几万张或者目标域和ImageNet差异很大比如医学影像、卫星图我建议全部微调但把学习率调低一些。常规的做法是整体学习率设0.0001左右分类头学习率可以稍微放大到0.001。ViT对学习率比较敏感一开始就用太大学习率很容易出现loss震荡甚至不收敛。另外一个小技巧是如果显存允许可以先把模型在较大分辨率如384x384上微调几轮效果通常会比224更好。因为ViT没有卷积的局部先验更大分辨率意味着更多patch更多细节。代价就是序列长度变长显存和速度都翻几倍。我自己做细粒度分类时这个技巧带来的精度提升很明显。4. 实操记录从安装到自定义数据集微调4.1 环境准备和安装先说我实测过的环境组合Python 3.10PyTorch 2.1以上CUDA 11.8显卡是RTX 3090。PyTorch的安装直接用官方命令就行如果下载慢可以把pip源换成清华源或者阿里源。装完之后装依赖pip install torch torchvision timm pip install vit-pytorch后面这个vit_pytorch就是lucidrains的版本也是很多博客里提到的vision-transformer-pytorch的PyTorch实现。不过我想特别提醒一句vit_pytorch这个包默认不携带预训练权重它的作用是快速构建模型结构你想直接跑预训练推理还是要配合timm或者transformers。很多人一开始没搞清楚这点装了个vit-pytorch然后发现模型是随机初始化的以为库有问题。4.2 快速推理用预训练权重分类一张图下面这段代码是我在项目里验证一张新图时常用的模板基于timm实现简单可靠import torch import timm from PIL import Image from torchvision import transforms device torch.device(cuda if torch.cuda.is_available() else cpu) model timm.create_model(vit_base_patch16_224, pretrainedTrue) model model.to(device) model.eval() transform transforms.Compose([ transforms.Resize((224, 224)), transforms.ToTensor(), transforms.Normalize(mean[0.485, 0.456, 0.406], std[0.229, 0.224, 0.225]), ]) img Image.open(test.jpg).convert(RGB) x transform(img).unsqueeze(0).to(device) with torch.no_grad(): out model(x) prob torch.softmax(out, dim1) top5 torch.topk(prob, 5)输出top5之后去查一下ImageNet的类别索引表就能知道模型预测的是什么类。这里最常踩的坑有两个一是忘了convert(RGB)导致灰度图或RGBA图报错二是忘了加batch维度unsqueeze(0)少写就报维度错误。我刚开始跑的时候就在这两处浪费过时间。4.3 自定义数据集微调假设现在你有一个10类的自定义数据集目录结构大概是train/class1、train/class2这样。用torchvision的ImageFolder读进来然后替换最后的分类头就可以开始微调import torch import torch.nn as nn import timm from torchvision import datasets, transforms from torch.utils.data import DataLoader model timm.create_model(vit_base_patch16_224, pretrainedTrue, num_classes10) model.to(device) train_transform transforms.Compose([ transforms.RandomResizedCrop(224), transforms.RandomHorizontalFlip(), transforms.ToTensor(), transforms.Normalize(mean[0.485, 0.456, 0.406], std[0.229, 0.224, 0.225]), ]) train_ds datasets.ImageFolder(train, transformtrain_transform) train_loader DataLoader(train_ds, batch_size32, shuffleTrue, num_workers4) optimizer torch.optim.AdamW(model.parameters(), lr1e-4, weight_decay0.05) criterion nn.CrossEntropyLoss() for epoch in range(10): model.train() for x, y in train_loader: x, y x.to(device), y.to(device) loss criterion(model(x), y) optimizer.zero_grad() loss.backward() optimizer.step() print(fepoch {epoch}, loss {loss.item():.4f})这段代码只是最基础的训练循环。实际项目里我还会加warmup、余弦退火、Mixup和数据增强。ViT在中小数据集上容易过拟合所以增强策略比CNN时代更讲究。还有一个细节AdamW的weight_decay我习惯设0.05这是从DeiT论文里来的经验值实测比0.01稳定。5. 避坑指南使用ViT时最容易踩的坑5.1 输入尺寸和归一化必须匹配ViT模型对输入尺寸非常敏感这一点比CNN严格得多。timm里的vit_base_patch16_224要求输入224x224如果你给它喂512x512的图patch数量就变了位置编码的维度对不上通常运行到模型内部就会直接报维度不匹配的错误。我的建议是把预处理统一封装成一个函数Resize到224归一化参数就用ImageNet默认的mean和std不要自己随便改。另外一旦换了数据集别忘了重新检查normalize参数是否匹配——尤其是医学图像或者红外图像它们的像素分布和自然图像差别很大。5.2 预训练权重下载失败怎么办这个问题在HuggingFace的权重上尤其常见。下载到一半断掉、网络超时、缓存损坏都会导致无法加载。我遇到这种情况一般分两步排查先看报错是不是SSL或者超时如果是就说明是网络问题可以设置HF_ENDPOINThttps://hf-mirror.com再重新拉取如果报错是键名不匹配那大概率是模型结构定义和权重来源版本不一致比如用了patch16的定义去加载patch32的权重。权重和模型结构的匹配相当重要。提示HuggingFace的缓存目录通常在~/.cache/huggingface删掉对应模型的缓存再重新下载可以解决很多奇怪的加载问题。下载慢但没报错时多试几次或者手动下载后放到缓存目录也行。5.3 显存不够用先别急着换显卡ViT虽然参数不算特别多但自注意力的计算复杂度是序列长度的平方。224分辨率下197个token还好一旦输入到448x448token数变成(448/16)^21785计算量增长非常明显。如果显存爆了最直接的思路是减小batch size或者把patch_size从16改成32。还有一个实用技巧是开启梯度累积用多个小batch累加梯度模拟大batch效果能接近但省显存。训练速度慢的话优先检查是不是数据加载瓶颈。num_workers调大或者用pin_memoryTrue经常能把GPU利用率拉满。我见过很多新人把num_workers默认0跑GPU利用率低得可怜改到4或8之后速度立竿见影。相比一上来就换卡这招划算得多。5.4 位置编码与输入分辨率不匹配如果你想在384x384或更大的分辨率上微调直接用vit_base_patch16_224的权重会报位置编码维度不匹配。因为预训练权重的position embedding是1x197x768而384x384对应的是1x577x768。解决办法是插值调整位置编码的尺寸。timm里有些版本支持img_size参数或者在create_model时指定img_size384但不是所有实现都自动处理。如果要手动插值可以这样import torch from vit_pytorch import ViT model ViT(...) pos_embed model.pos_embedding # 形状 [1, 197, 768] new_pos_embed torch.nn.functional.interpolate( pos_embed.permute(0, 2, 1).unsqueeze(0), size(577,), modelinear, align_cornersFalse, ).squeeze(0).permute(0, 2, 1) model.pos_embedding torch.nn.Parameter(new_pos_embed)不过说实话如果只是做普通分类任务我不建议手动插值直接用timm里带384后缀的模型比如vit_base_patch16_384会省心很多位置编码部分timm已经处理好了。把上面这些坑都踩过一遍之后我对ViT的理解反而更深了。说实话ViT的门槛并不在模型本身而在于各种细节patch怎么切、位置编码怎么处理、预训练权重怎么融合、微调参数怎么调。把这些细节弄明白它就是非常趁手的视觉骨干网络。如果你刚开始接触ViT我的建议是先按着代码把模型结构打印出来一个个模块核对再跑通预训练推理最后再上自己的数据。这个过程走一遍比干啃论文有用得多。最后分享一个我自己的习惯任何新项目要采用ViT我都会先用Base模型加ImageNet预训练跑一个baseline确认任务可行之后再根据显存和精度需求决定要不要换Large或者调分辨率。不要一开始就上大模型否则调参和排错的成本会高到让你怀疑人生。希望这篇内容对你有帮助有问题欢迎留言交流。本文还有配套的精品资源点击获取