说实话第一次看到“ViT做图像分类”这个标题时我心里是犯过嘀咕的。毕竟做了好几年卷积网络突然让我把图像切成一小块一小块然后塞给一套NLP里常用的Transformer架构总觉得是某种“跨界炫技”。但真的动手跑通一次之后我的态度彻底改变了。ViTVision Transformer这套思路不仅把图像分类从“卷积一统天下”的局面里撕开了一个口子而且从迁移学习、可视化分析到后续的多模态扩展它都成了绕不开的基础底座。这篇文章我不会给你空谈原理再丢一堆论文链接。我会按照我从零搭建、训练、调优到最终可视化的完整流程来写尽量把每一步的“为什么这样做”也讲清楚。如果你是第一次接触ViT或者已经跑过一些分类模型、想换个架构试试水这篇文章应该能直接帮你省掉不少查资料的功夫。目标是从环境准备开始到最终拿到训练曲线、混淆矩阵和注意力热力图一条链路走完不卡壳。1. 先把ViT的“身位”看清楚它到底解决什么问题在动手敲代码之前我想先花一点时间把ViT的价值定位说清楚。很多人一上来就急着装库、跑模型结果代码能跑但心里始终是模糊的——这会影响后续改数据集、调参数时的判断力。1.1 卷积网络靠“局部归纳偏置”ViT靠“全局注意力”传统卷积网络比如ResNet系列之所以在图像任务上表现好是因为它天生带着两个很强的先验局部性和平移等变性。也就是说卷积假设图像里相邻的像素关系最密切而且物体出现在哪里不太影响识别结果。这两个先验让卷积网络在小数据集上也能学得不错因为它要学的东西被“约束”住了。ViT的思路则完全不同。它把一张图切分成固定大小的patch比如16x16像素一块每个patch展平后当成一个“词”再把这些“词”的组合关系交给Transformer的注意力机制去建模。注意力机制没有局部性假设它允许任意两个位置直接建立关联。听起来很美好但代价也很实在纯粹从零训练一个ViT效果往往打不过同量级的ResNet。原因很简单通用性强的模型需要更多数据去“喂”出那些卷积网络天生自带的知识。ViT真正发力的场景是大规模预训练加下游微调——这就是为什么你在实际项目里用ViT基本都是加载预训练权重起步。1.2 什么场景适合用它根据自己的实际项目经验我一般这样判断数据量中等以上或者能借助大规模预训练模型做微调适合用ViT效果通常能摸到当时的第一梯队。想要可视化模型到底“看”了图像哪些区域ViT的注意力图天然适合做这件事后面我会给出具体方法。从卷积切到Transformer为后续做多模态、跨模态任务打基础ViT是很好的第一步跳板。训练数据极少几千张以内又没有合适的预训练权重不建议硬上ViT老老实实用卷积或者先用小模型蒸馏。1.3 用一张表记住ViT的几个关键变体实际使用中你大概率会遇到这几个配置我先列成一张速查表。不同规模的模型对应不同的参数量和显存需求看完你就知道选型时该往哪边靠模型配置patch size参数量级常用输入分辨率适用场景ViT-Tiny16约500万224x224快速实验、显存受限ViT-Small16约2200万224x224中等任务微调ViT-Base16约8600万224x224 / 384x384通用分类、迁移学习主力ViT-Large16约3亿224x224 / 384x384大数据集、高精度需求最常用的其实是ViT-Base搭配16x16的patch这个组合在精度和训练成本之间最平衡。后面我的代码也是基于这个配置做的。2. 环境搭建的最小闭环我花了不到十分钟标题说“5分钟上手”严格来说有点夸张但如果只是搭建一个能跑ViT训练的环境十分钟确实够了。我把我的搭建过程拆开写你可以照着做每一步都验证过。2.1 创建干净的Python环境我习惯先用conda隔离环境避免把系统Python搞乱。conda create -n vit_env python3.10 -y conda activate vit_envPython版本选3.10是因为PyTorch和各个库的兼容性最稳。3.11、3.12也不是不行但没必要在环境上给自己找麻烦。2.2 安装PyTorch和需要用到的库PyTorch的安装命令会根据你的机器是否支持CUDA而不同。我这里给一个通用做法如果你有NVIDIA显卡优先去PyTorch官网选择对应CUDA版本的安装命令如果是纯CPU环境直接装CPU版也能跑通流程只是训练慢一些。# CPU版本保证能跑通 pip install torch torchvision # 如果用的是CUDA去官网生成对应命令即可然后安装ViT相关模型库和可视化工具pip install timm matplotlib scikit-learn seaborn tensorboard解释一下这几个库的用途timm一个开源的图像模型库里面集成了ViT、Swin、ConvNeXt等各种模型而且提供了预训练权重。你不用自己从零写ViT网络结构直接调模型名就行。matplotlib / seaborn画训练曲线和混淆矩阵用。scikit-learn计算分类报告、混淆矩阵用。tensorboard记录训练过程中的loss和acc曲线方便长时间训练时观察状态。2.3 快速验证环境是否就绪安装完成后先写一行代码验证模型能不能加载import timm model timm.create_model(vit_base_patch16_224, pretrainedTrue, num_classes10) print(model)如果最后一层输出是10类说明模型结构已经替换好了。这一步总共不到两分钟。看到模型打印出来的那一刻环境问题就基本结束了。这里有个小经验timm里模型名称的规律是“架构名_patch尺寸_输入分辨率”比如vit_base_patch16_224就是ViT-Base、patch16、输入224x224。命名规则记住以后你想要什么配置都在这个框架里选不用背参数。3. 核心机制拆解图像是怎么变成一串Token的环境装好之后别急着开训。我建议你先花十分钟把ViT的数据流走一遍因为后面调试错误时90%的问题都出在对维度变化的理解上。3.1 Patch Embedding把图像切成“单词”假设输入是一张3通道、224x224的彩色图。ViT的第一步是把它切成16x16的小块。224除以16等于14所以图像会被切成14x14196个小块。每一个16x16x3的小块展平后是一个长度为768的向量16乘16乘3。这196个向量就是Transformer眼里的“词序列”。序列长度是196每个词的维度是768。这一步叫Patch Embedding本质上就是一次卷积加展平的操作它完成了从像素空间到向量空间的转换。你可以这样类比一张图片原本是一整块信息ViT把它当成一篇由196个“单词”组成的短文每个“单词”描述局部区域的颜色和纹理。注意力机制负责理解这些“单词”之间的关系。3.2 位置编码和CLS Token让模型知道顺序与全局信息Transformer的注意力机制本身是不关心顺序的它把输入当成“袋子里的词”。但图像patch的顺序显然有意义——左上角的天空和右下角的地面不能随意互换。所以ViT要加一个位置编码向量把每个patch的位置信息叠加到向量里。ViT用的是可学习的位置编码训练时它会自动学到“第i个位置大概长什么样”。另外ViT借鉴了BERT的做法在序列开头加一个特殊的CLS Token也是768维。Transformer编码器跑完整个序列后我们只取CLS Token对应的输出作为整张图的总表征再接分类头。为什么不用所有patch的均值因为CLS Token在注意力机制里可以自由地和所有patch交互它有机会聚合全局信息而均值池化其实也能用只是从实践效果看CLS Token微调时往往更稳。3.3 Transformer Encoder信息在这里交互196个patch向量加上一个CLS Token一共197个Token一起进入Transformer Encoder。Encoder的核心是多头自注意力Multi-Head Self-Attention和前馈网络MLP中间还带着残差连接和LayerNorm。多头注意力的作用可以简单理解为让每个Token在多个“视角”下查看其他Token。某个视角可能关注纹理相似性另一个视角可能关注空间相邻性还有一个视角可能关注语义相关性。多个视角的结果拼接起来就形成了一次全局信息交换。实际代码中你并不需要手动实现这个模块timm已经帮你封装好了。但你必须知道197个Token是固定长度所以输入图像分辨率变了patch数量就会变网络的序列长度跟着变这直接影响位置编码是否匹配。3.4 一张图看懂维度流转我习惯用维度表来梳理流程这样比自己死记公式快阶段张量维度以ViT-Base为例输入图像[B, 3, 224, 224]Patch切分后嵌入[B, 196, 768]拼接CLS Token[B, 197, 768]加入位置编码[B, 197, 768]Transformer Encoder输出[B, 197, 768]取CLS Token映射到类别[B, num_classes]B是batch size。只要确认每一步的维度都对得上代码基本就不会出大问题。4. 数据集准备与预处理从CIFAR-10到自己的图片集原理过完我们回到实操。图像分类项目里最耗时、最容易出错的往往不是模型而是数据准备。我先用CIFAR-10做示例因为内置下载、标注齐全适合快速验证训练链路。等模型跑通了再替换成你自己的数据集逻辑是一样的。4.1 用内置数据集跑通全流程CIFAR-10是10类32x32的小图集虽然分辨率低但作为流程验证非常合适。要注意一个问题ViT-Base的预训练输入是224x224直接把32x32的小图塞进去会丢失大量细节。所以工程上通常把CIFAR-10放大到224x224再输入。from torchvision import datasets, transforms train_transform transforms.Compose([ transforms.Resize((224, 224)), transforms.RandomHorizontalFlip(), transforms.ToTensor(), transforms.Normalize(mean[0.5, 0.5, 0.5], std[0.5, 0.5, 0.5]), ]) val_transform transforms.Compose([ transforms.Resize((224, 224)), transforms.ToTensor(), transforms.Normalize(mean[0.5, 0.5, 0.5], std[0.5, 0.5, 0.5]), ]) train_set datasets.CIFAR10(root./data, trainTrue, downloadTrue, transformtrain_transform) val_set datasets.CIFAR10(root./data, trainFalse, downloadTrue, transformval_transform)随机水平翻转是轻量级数据增强能让模型泛化稍好一些。后面如果你想冲更高精度还可以加RandomResizedCrop、MixUp等更强的手段但基线阶段先保持简单。4.2 替换成自己的数据集目录当你想训练自己的图片分类任务时最省事的方式是ImageFolder格式的目录结构dataset/ train/ class_a/ 001.jpg 002.jpg class_b/ 001.jpg val/ class_a/ 001.jpg class_b/ 001.jpg然后用两行代码加载from torchvision import datasets train_set datasets.ImageFolder(rootdataset/train, transformtrain_transform) val_set datasets.ImageFolder(rootdataset/val, transformval_transform) print(train_set.classes) # 自动读取类别名称只要目录结构对了类别数量、类别名称都不用手动写非常省事。唯一要注意的是每个类别下的图片数量尽量均衡如果某个类别只有十几张训练效果会明显偏向样本多的类。4.3 DataLoader的batch size怎么定这一步直接关系到训练速度和显存占用。ViT-Base在224分辨率下单张图的前向显存占用不低。我自己的经验是10G左右显存batch size设16或32。6G左右显存batch size设8或16。显存不够时首选减小batch size其次才是换小模型。from torch.utils.data import DataLoader train_loader DataLoader(train_set, batch_size32, shuffleTrue, num_workers4) val_loader DataLoader(val_set, batch_size32, shuffleFalse, num_workers4)num_workers设为4是为了让数据加载和模型训练并行不至于GPU等着CPU传数据。Windows环境下num_workers如果报错可以改为0但训练速度会受影响。5. 训练与验证完整代码与调参思路数据准备好了模型加载好了接下来把训练链路打通。这一章是整个项目的核心我会先给出一个能直接跑的完整训练循环然后逐个解释关键选择。5.1 加载预训练模型并替换分类头import torch import torch.nn as nn import timm model timm.create_model(vit_base_patch16_224, pretrainedTrue, num_classes10) model.cuda()这里两个关键参数pretrainedTrue加载在ImageNet上预训练的权重。这是ViT能碾压小数据集的关键千万不要省略。num_classes10timm会自动把模型原来自带的1000类分类头替换成10类。这一步相当于只保留前面积累的视觉特征把最后的判断层换成你自己的任务。5.2 损失函数、优化器和学习率策略criterion nn.CrossEntropyLoss() optimizer torch.optim.AdamW(model.parameters(), lr3e-5, weight_decay0.05)损失函数用交叉熵这是分类任务的标配。优化器选AdamW因为它在Transformer架构上表现稳定weight_decay设置0.05是ViT训练中比较常见的正则化做法。学习率这里要特别强调从预训练模型微调时学习率一定不能给太大。3e-5是个安全起点ViT对学习率比卷积网络更敏感给的太大很容易出现loss振荡甚至直接发散。我个人习惯再加一个warmup策略前几个epoch用很小的学习率“热热身”再逐步升到目标学习率。这样可以避免预训练权重在训练一开始就被破坏。from torch.optim.lr_scheduler import CosineAnnealingLR scheduler CosineAnnealingLR(optimizer, T_max20, eta_min1e-6)余弦退火是ViT训练里很稳妥的调度方式。前期学习率稍高后期平滑下降最后收敛得更干净。5.3 训练循环的完整框架def train_one_epoch(model, loader, optimizer, criterion): model.train() total_loss, correct, total 0, 0, 0 for x, y in loader: x, y x.cuda(), y.cuda() optimizer.zero_grad() logits model(x) loss criterion(logits, y) loss.backward() optimizer.step() total_loss loss.item() * x.size(0) correct (logits.argmax(dim1) y).sum().item() total x.size(0) return total_loss / total, correct / total torch.no_grad() def validate(model, loader, criterion): model.eval() total_loss, correct, total 0, 0, 0 for x, y in loader: x, y x.cuda(), y.cuda() logits model(x) loss criterion(logits, y) total_loss loss.item() * x.size(0) correct (logits.argmax(dim1) y).sum().item() total x.size(0) return total_loss / total, correct / total训练和验证的区别就是是否计算梯度、是否更新参数。注意验证函数上加了torch.no_grad()这样会省掉大量中间变量缓存既省显存又加快速度。主循环epochs 20 best_acc 0 for epoch in range(epochs): train_loss, train_acc train_one_epoch(model, train_loader, optimizer, criterion) val_loss, val_acc validate(model, val_loader, criterion) scheduler.step() print(fEpoch {epoch1}: train_loss{train_loss:.4f}, train_acc{train_acc:.4f}, val_acc{val_acc:.4f}) if val_acc best_acc: best_acc val_acc torch.save(model.state_dict(), best_model.pth)我只跑了20个epoch做基线验证。CIFAR-10这个任务上ViT-Base微调20轮之后验证精度大概能到90%以上如果你把训练策略调得更精细还能更高。5.4 训练中观察到的现象我在实际跑的过程中有几个体会前两三个epoch准确率会快速上升这是预训练权重在下游任务上“热身”的过程。如果loss一直不降先检查学习率是不是太大或者预处理中Normalize的均值方差是不是写错了。验证集准确率偶尔会抖动不用慌只要整体趋势向上就行最终模型以验证集最优为准。6. 可视化输出训练曲线、混淆矩阵与注意力热力图模型训练完光有一个准确率数字远远不够。你需要把训练过程、分类结果、模型关注区域都画出来。这个章节也是“从环境搭建到可视化输出”的最后一环。6.1 绘制训练曲线训练曲线能直观反映收敛状态。最简单的方式是用matplotlib记录每个epoch的train_loss和val_accimport matplotlib.pyplot as plt # history {train_loss: [], val_acc: []} plt.figure(figsize(12, 4)) plt.subplot(1, 2, 1) plt.plot(history[train_loss], labeltrain_loss) plt.xlabel(epoch) plt.ylabel(loss) plt.legend() plt.subplot(1, 2, 2) plt.plot(history[val_acc], labelval_acc) plt.xlabel(epoch) plt.ylabel(accuracy) plt.legend() plt.tight_layout() plt.savefig(training_curves.png)如果你习惯用TensorBoard也可以在训练循环里加入SummaryWriter记录。TensorBoard的优势是能实时看曲线不用等训练结束。6.2 绘制混淆矩阵混淆矩阵能告诉你模型到底把哪些类搞混了。我用scikit-learn计算再用seaborn画图。import numpy as np from sklearn.metrics import confusion_matrix import seaborn as sns all_preds, all_labels [], [] model.eval() with torch.no_grad(): for x, y in val_loader: x x.cuda() logits model(x) preds logits.argmax(dim1).cpu().numpy() all_preds.extend(preds) all_labels.extend(y.numpy()) cm confusion_matrix(all_labels, all_preds) plt.figure(figsize(8, 8)) sns.heatmap(cm, annotTrue, fmtd, cmapBlues) plt.xlabel(Predicted) plt.ylabel(True) plt.savefig(confusion_matrix.png)如果发现某两个类互相混淆严重通常说明它们在视觉上确实相似或者样本数量不均。这时候可以针对性增加数据增强或者补充这类图片的训练量。6.3 注意力热力图模型到底在看哪里这是ViT最吸引人的地方。因为Transformer的注意力权重天然可以提取出来告诉你模型做判断时重点参考了哪些patch。你可以用下面的思路获取最后一层Encoder的注意力图import torch.nn.functional as F # 注册一个钩子获取最后一个Transformer Block的注意力权重 attention_maps [] def hook_fn(module, input, output): # output 里包含 attention 权重需要按实际模型结构取 attention_maps.append(output) # 以timm的ViT为例通常可以这样访问 model.blocks[-1].attn.register_forward_hook(hook_fn)更务实的做法是直接使用timm内部提供的get_intermediate_layers方法它能返回每个block的输出特征。取最后一层所有注意力头的平均值按patch的14x14网格还原成空间尺寸再放大到原图大小。with torch.no_grad(): # 输入一张(1,3,224,224)的图 feats model.get_intermediate_layers(x, n1)[0] # [1, 197, 768] # 得到的是最后一层的特征不是注意力权重如果你只是想快速看到模型关注区域更实用的方案是使用GradCAM类方法作用于ViT的最后一个卷积层或注意力块。对于ViT可以基于最后一层所有Token对CLS Token的注意力权重来生成热力图。基本思路是取出最后一层注意力权重中CLS Token对应的那一行丢掉CLS自己剩下的196个值reshape成14x14上采样回224x224再与原图叠加。这个热力图能非常直观地告诉你模型是否关注了目标区域。我在实际项目里发现模型分类正确时热力图通常集中在主体物体上如果热力图乱飘那基本可以断定模型在“抄近道”靠背景信息作弊这时候就要考虑是否数据有泄漏或者增强方式是否需要调整。7. 我踩过的几个坑ViT训练中反直觉的细节最后这一节不写教程式的标准步骤纯粹分享我实际踩过、也花了不少时间排查的坑。每一条都应该能帮你在自己的项目中少走弯路。7.1 预训练模型的输入尺寸必须对齐第一次用ViT时我直接拿256x256的图像输入心想“分辨率大一点信息多一点总没坏处”。结果模型直接报维度错误。原因在于位置编码的数量是训练时定死的224分辨率对应196个patch256分辨率则对应256个patch位置编码矩阵对不上。timm里如果你传pretrainedTrue它会用224对应的编码。想用大图就得用支持可变分辨率的位置编码插值方案或者干脆选针对384分辨率预训练的变体。我后来的习惯是除非特别需要否则就用224x224省心又够用。7.2 学习率给大了一轮就崩卷积网络微调时我用1e-4经常没事换到ViT之后第一次也照抄了这个值结果训练到第5轮loss开始往上飙。后来检查发现ViT对学习率的容忍度明显更低。把学习率降到3e-5之后训练一下就稳了。我的经验值是ViT微调学习率从1e-5到5e-5之间起步比较安全。你要是实在想冲更快也建议先用3e-5跑通一轮确认loss下降稳定再加。7.3 忘了切换训练/验证模式这个错误特别低幼但真的容易忽略。验证阶段如果不调用model.eval()BatchNorm和Dropout的行为会和训练时不一样准确率可能忽高忽低。ViT里虽然BatchNorm用得少但Dropout依然存在影响不小。7.4 显存溢出不一定是batch size的锅有一次我训练中途报显存溢出第一反应是减小batch size结果发现是验证阶段忘记写torch.no_grad()导致推理时也保留了完整计算图。这个坑很隐蔽因为训练阶段显存本来就用得多溢出时不容易想到是验证代码的问题。7.5 自己数据集类别数和预训练头不一致别忽略很多人觉得timm里num_classes参数传一下就行其实替换分类头之后新初始化的权重是随机值前几轮梯度主要集中在这一层。如果你的数据集很小新分类头的训练可能不够充分。我的解决办法是给新分类头单独设一个稍大的学习率或者先把分类头单独训练几个epoch再全量微调。这个小技巧在数据量少的时候往往会带来意外惊喜。8. 从基线到进阶后续还能怎么扩展到这里你已经跑通了一条完整的ViT图像分类链路也拿到了可视化结果。但我还想多说几句后续的扩展方向因为ViT的想象力远不止“替掉ResNet做分类”这么简单。8.1 把注意力热力图做成线上诊断工具我在真实项目中会把注意力热力图写入TensorBoard每轮训练结束随机抽几张验证图看模型的注意力是否稳定落在目标上。这个操作成本极低但对模型调优帮助很大。你会发现有些时候模型“准确率高”其实是一种假象它靠的是背景或水印。这类问题只有可视化才能看出来。8.2 尝试更强的数据增强策略如果你想在CIFAR-10或类似小数据集上冲击更高精度可以在预处理里加入RandomResizedCrop、RandAugment、MixUp和CutMix。我自己实验下来这些增强策略对ViT的帮助比传统卷积网络更明显。原因是ViT的归纳偏置少需要靠数据增强“补课”的地方更多。8.3 从分类到其他任务ViT的分类头拿掉之后它就是一个通用的视觉特征提取器。你可以在此基础上做图像检索、零样本分类甚至接入文本编码器做图文匹配。如果你之后要接触多模态模型ViT的架构知识几乎是必需品。9. 一些个人经验与最后的建议折腾ViT的这段时间我最深的体会是ViT本身不复杂复杂的是你脑子里的那套“卷积思维”要切换过来。刚开始我老是下意识地问“感受野多大”“池化在哪”但这些概念在ViT里都被重新定义了。放下这些旧框架之后代码写起来反而更简单因为整个流程就是“切块、编码、Transformer、接分类头”。如果你准备动手实践我的建议是别急于一次跑通所有内容。按照最简单的路径来先加载预训练模型在CIFAR-10上跑通一个epoch确认训练曲线在下降接着加上完整的验证和模型保存最后再尝试可视化。一步一步来出了问题时定位会非常清晰。最后分享一个很实用的小技巧加载完预训练权重后别急着改任何训练参数先用一张图片做一次前向推理直接看图能不能输出合理的类别概率。如果这一步就错了说明环境或者模型加载有问题后面怎么调都白搭。我的所有失败项目几乎都在这类“前置检查没做”的情况下浪费过时间。希望这篇文章能让你少走一些弯路。如果跑的过程中遇到什么奇怪的报错欢迎回来对照着第七章的坑查一遍大概率问题就出在那里。 
长河君 · 2026/10/20 2:58:50 · 阅读 1.2k ViT图像分类实战教程：从环境搭建到可视化全流程
说实话第一次看到“ViT做图像分类”这个标题时我心里是犯过嘀咕的。毕竟做了好几年卷积网络突然让我把图像切成一小块一小块然后塞给一套NLP里常用的Transformer架构总觉得是某种“跨界炫技”。但真的动手跑通一次之后我的态度彻底改变了。ViTVision Transformer这套思路不仅把图像分类从“卷积一统天下”的局面里撕开了一个口子而且从迁移学习、可视化分析到后续的多模态扩展它都成了绕不开的基础底座。这篇文章我不会给你空谈原理再丢一堆论文链接。我会按照我从零搭建、训练、调优到最终可视化的完整流程来写尽量把每一步的“为什么这样做”也讲清楚。如果你是第一次接触ViT或者已经跑过一些分类模型、想换个架构试试水这篇文章应该能直接帮你省掉不少查资料的功夫。目标是从环境准备开始到最终拿到训练曲线、混淆矩阵和注意力热力图一条链路走完不卡壳。1. 先把ViT的“身位”看清楚它到底解决什么问题在动手敲代码之前我想先花一点时间把ViT的价值定位说清楚。很多人一上来就急着装库、跑模型结果代码能跑但心里始终是模糊的——这会影响后续改数据集、调参数时的判断力。1.1 卷积网络靠“局部归纳偏置”ViT靠“全局注意力”传统卷积网络比如ResNet系列之所以在图像任务上表现好是因为它天生带着两个很强的先验局部性和平移等变性。也就是说卷积假设图像里相邻的像素关系最密切而且物体出现在哪里不太影响识别结果。这两个先验让卷积网络在小数据集上也能学得不错因为它要学的东西被“约束”住了。ViT的思路则完全不同。它把一张图切分成固定大小的patch比如16x16像素一块每个patch展平后当成一个“词”再把这些“词”的组合关系交给Transformer的注意力机制去建模。注意力机制没有局部性假设它允许任意两个位置直接建立关联。听起来很美好但代价也很实在纯粹从零训练一个ViT效果往往打不过同量级的ResNet。原因很简单通用性强的模型需要更多数据去“喂”出那些卷积网络天生自带的知识。ViT真正发力的场景是大规模预训练加下游微调——这就是为什么你在实际项目里用ViT基本都是加载预训练权重起步。1.2 什么场景适合用它根据自己的实际项目经验我一般这样判断数据量中等以上或者能借助大规模预训练模型做微调适合用ViT效果通常能摸到当时的第一梯队。想要可视化模型到底“看”了图像哪些区域ViT的注意力图天然适合做这件事后面我会给出具体方法。从卷积切到Transformer为后续做多模态、跨模态任务打基础ViT是很好的第一步跳板。训练数据极少几千张以内又没有合适的预训练权重不建议硬上ViT老老实实用卷积或者先用小模型蒸馏。1.3 用一张表记住ViT的几个关键变体实际使用中你大概率会遇到这几个配置我先列成一张速查表。不同规模的模型对应不同的参数量和显存需求看完你就知道选型时该往哪边靠模型配置patch size参数量级常用输入分辨率适用场景ViT-Tiny16约500万224x224快速实验、显存受限ViT-Small16约2200万224x224中等任务微调ViT-Base16约8600万224x224 / 384x384通用分类、迁移学习主力ViT-Large16约3亿224x224 / 384x384大数据集、高精度需求最常用的其实是ViT-Base搭配16x16的patch这个组合在精度和训练成本之间最平衡。后面我的代码也是基于这个配置做的。2. 环境搭建的最小闭环我花了不到十分钟标题说“5分钟上手”严格来说有点夸张但如果只是搭建一个能跑ViT训练的环境十分钟确实够了。我把我的搭建过程拆开写你可以照着做每一步都验证过。2.1 创建干净的Python环境我习惯先用conda隔离环境避免把系统Python搞乱。conda create -n vit_env python3.10 -y conda activate vit_envPython版本选3.10是因为PyTorch和各个库的兼容性最稳。3.11、3.12也不是不行但没必要在环境上给自己找麻烦。2.2 安装PyTorch和需要用到的库PyTorch的安装命令会根据你的机器是否支持CUDA而不同。我这里给一个通用做法如果你有NVIDIA显卡优先去PyTorch官网选择对应CUDA版本的安装命令如果是纯CPU环境直接装CPU版也能跑通流程只是训练慢一些。# CPU版本保证能跑通 pip install torch torchvision # 如果用的是CUDA去官网生成对应命令即可然后安装ViT相关模型库和可视化工具pip install timm matplotlib scikit-learn seaborn tensorboard解释一下这几个库的用途timm一个开源的图像模型库里面集成了ViT、Swin、ConvNeXt等各种模型而且提供了预训练权重。你不用自己从零写ViT网络结构直接调模型名就行。matplotlib / seaborn画训练曲线和混淆矩阵用。scikit-learn计算分类报告、混淆矩阵用。tensorboard记录训练过程中的loss和acc曲线方便长时间训练时观察状态。2.3 快速验证环境是否就绪安装完成后先写一行代码验证模型能不能加载import timm model timm.create_model(vit_base_patch16_224, pretrainedTrue, num_classes10) print(model)如果最后一层输出是10类说明模型结构已经替换好了。这一步总共不到两分钟。看到模型打印出来的那一刻环境问题就基本结束了。这里有个小经验timm里模型名称的规律是“架构名_patch尺寸_输入分辨率”比如vit_base_patch16_224就是ViT-Base、patch16、输入224x224。命名规则记住以后你想要什么配置都在这个框架里选不用背参数。3. 核心机制拆解图像是怎么变成一串Token的环境装好之后别急着开训。我建议你先花十分钟把ViT的数据流走一遍因为后面调试错误时90%的问题都出在对维度变化的理解上。3.1 Patch Embedding把图像切成“单词”假设输入是一张3通道、224x224的彩色图。ViT的第一步是把它切成16x16的小块。224除以16等于14所以图像会被切成14x14196个小块。每一个16x16x3的小块展平后是一个长度为768的向量16乘16乘3。这196个向量就是Transformer眼里的“词序列”。序列长度是196每个词的维度是768。这一步叫Patch Embedding本质上就是一次卷积加展平的操作它完成了从像素空间到向量空间的转换。你可以这样类比一张图片原本是一整块信息ViT把它当成一篇由196个“单词”组成的短文每个“单词”描述局部区域的颜色和纹理。注意力机制负责理解这些“单词”之间的关系。3.2 位置编码和CLS Token让模型知道顺序与全局信息Transformer的注意力机制本身是不关心顺序的它把输入当成“袋子里的词”。但图像patch的顺序显然有意义——左上角的天空和右下角的地面不能随意互换。所以ViT要加一个位置编码向量把每个patch的位置信息叠加到向量里。ViT用的是可学习的位置编码训练时它会自动学到“第i个位置大概长什么样”。另外ViT借鉴了BERT的做法在序列开头加一个特殊的CLS Token也是768维。Transformer编码器跑完整个序列后我们只取CLS Token对应的输出作为整张图的总表征再接分类头。为什么不用所有patch的均值因为CLS Token在注意力机制里可以自由地和所有patch交互它有机会聚合全局信息而均值池化其实也能用只是从实践效果看CLS Token微调时往往更稳。3.3 Transformer Encoder信息在这里交互196个patch向量加上一个CLS Token一共197个Token一起进入Transformer Encoder。Encoder的核心是多头自注意力Multi-Head Self-Attention和前馈网络MLP中间还带着残差连接和LayerNorm。多头注意力的作用可以简单理解为让每个Token在多个“视角”下查看其他Token。某个视角可能关注纹理相似性另一个视角可能关注空间相邻性还有一个视角可能关注语义相关性。多个视角的结果拼接起来就形成了一次全局信息交换。实际代码中你并不需要手动实现这个模块timm已经帮你封装好了。但你必须知道197个Token是固定长度所以输入图像分辨率变了patch数量就会变网络的序列长度跟着变这直接影响位置编码是否匹配。3.4 一张图看懂维度流转我习惯用维度表来梳理流程这样比自己死记公式快阶段张量维度以ViT-Base为例输入图像[B, 3, 224, 224]Patch切分后嵌入[B, 196, 768]拼接CLS Token[B, 197, 768]加入位置编码[B, 197, 768]Transformer Encoder输出[B, 197, 768]取CLS Token映射到类别[B, num_classes]B是batch size。只要确认每一步的维度都对得上代码基本就不会出大问题。4. 数据集准备与预处理从CIFAR-10到自己的图片集原理过完我们回到实操。图像分类项目里最耗时、最容易出错的往往不是模型而是数据准备。我先用CIFAR-10做示例因为内置下载、标注齐全适合快速验证训练链路。等模型跑通了再替换成你自己的数据集逻辑是一样的。4.1 用内置数据集跑通全流程CIFAR-10是10类32x32的小图集虽然分辨率低但作为流程验证非常合适。要注意一个问题ViT-Base的预训练输入是224x224直接把32x32的小图塞进去会丢失大量细节。所以工程上通常把CIFAR-10放大到224x224再输入。from torchvision import datasets, transforms train_transform transforms.Compose([ transforms.Resize((224, 224)), transforms.RandomHorizontalFlip(), transforms.ToTensor(), transforms.Normalize(mean[0.5, 0.5, 0.5], std[0.5, 0.5, 0.5]), ]) val_transform transforms.Compose([ transforms.Resize((224, 224)), transforms.ToTensor(), transforms.Normalize(mean[0.5, 0.5, 0.5], std[0.5, 0.5, 0.5]), ]) train_set datasets.CIFAR10(root./data, trainTrue, downloadTrue, transformtrain_transform) val_set datasets.CIFAR10(root./data, trainFalse, downloadTrue, transformval_transform)随机水平翻转是轻量级数据增强能让模型泛化稍好一些。后面如果你想冲更高精度还可以加RandomResizedCrop、MixUp等更强的手段但基线阶段先保持简单。4.2 替换成自己的数据集目录当你想训练自己的图片分类任务时最省事的方式是ImageFolder格式的目录结构dataset/ train/ class_a/ 001.jpg 002.jpg class_b/ 001.jpg val/ class_a/ 001.jpg class_b/ 001.jpg然后用两行代码加载from torchvision import datasets train_set datasets.ImageFolder(rootdataset/train, transformtrain_transform) val_set datasets.ImageFolder(rootdataset/val, transformval_transform) print(train_set.classes) # 自动读取类别名称只要目录结构对了类别数量、类别名称都不用手动写非常省事。唯一要注意的是每个类别下的图片数量尽量均衡如果某个类别只有十几张训练效果会明显偏向样本多的类。4.3 DataLoader的batch size怎么定这一步直接关系到训练速度和显存占用。ViT-Base在224分辨率下单张图的前向显存占用不低。我自己的经验是10G左右显存batch size设16或32。6G左右显存batch size设8或16。显存不够时首选减小batch size其次才是换小模型。from torch.utils.data import DataLoader train_loader DataLoader(train_set, batch_size32, shuffleTrue, num_workers4) val_loader DataLoader(val_set, batch_size32, shuffleFalse, num_workers4)num_workers设为4是为了让数据加载和模型训练并行不至于GPU等着CPU传数据。Windows环境下num_workers如果报错可以改为0但训练速度会受影响。5. 训练与验证完整代码与调参思路数据准备好了模型加载好了接下来把训练链路打通。这一章是整个项目的核心我会先给出一个能直接跑的完整训练循环然后逐个解释关键选择。5.1 加载预训练模型并替换分类头import torch import torch.nn as nn import timm model timm.create_model(vit_base_patch16_224, pretrainedTrue, num_classes10) model.cuda()这里两个关键参数pretrainedTrue加载在ImageNet上预训练的权重。这是ViT能碾压小数据集的关键千万不要省略。num_classes10timm会自动把模型原来自带的1000类分类头替换成10类。这一步相当于只保留前面积累的视觉特征把最后的判断层换成你自己的任务。5.2 损失函数、优化器和学习率策略criterion nn.CrossEntropyLoss() optimizer torch.optim.AdamW(model.parameters(), lr3e-5, weight_decay0.05)损失函数用交叉熵这是分类任务的标配。优化器选AdamW因为它在Transformer架构上表现稳定weight_decay设置0.05是ViT训练中比较常见的正则化做法。学习率这里要特别强调从预训练模型微调时学习率一定不能给太大。3e-5是个安全起点ViT对学习率比卷积网络更敏感给的太大很容易出现loss振荡甚至直接发散。我个人习惯再加一个warmup策略前几个epoch用很小的学习率“热热身”再逐步升到目标学习率。这样可以避免预训练权重在训练一开始就被破坏。from torch.optim.lr_scheduler import CosineAnnealingLR scheduler CosineAnnealingLR(optimizer, T_max20, eta_min1e-6)余弦退火是ViT训练里很稳妥的调度方式。前期学习率稍高后期平滑下降最后收敛得更干净。5.3 训练循环的完整框架def train_one_epoch(model, loader, optimizer, criterion): model.train() total_loss, correct, total 0, 0, 0 for x, y in loader: x, y x.cuda(), y.cuda() optimizer.zero_grad() logits model(x) loss criterion(logits, y) loss.backward() optimizer.step() total_loss loss.item() * x.size(0) correct (logits.argmax(dim1) y).sum().item() total x.size(0) return total_loss / total, correct / total torch.no_grad() def validate(model, loader, criterion): model.eval() total_loss, correct, total 0, 0, 0 for x, y in loader: x, y x.cuda(), y.cuda() logits model(x) loss criterion(logits, y) total_loss loss.item() * x.size(0) correct (logits.argmax(dim1) y).sum().item() total x.size(0) return total_loss / total, correct / total训练和验证的区别就是是否计算梯度、是否更新参数。注意验证函数上加了torch.no_grad()这样会省掉大量中间变量缓存既省显存又加快速度。主循环epochs 20 best_acc 0 for epoch in range(epochs): train_loss, train_acc train_one_epoch(model, train_loader, optimizer, criterion) val_loss, val_acc validate(model, val_loader, criterion) scheduler.step() print(fEpoch {epoch1}: train_loss{train_loss:.4f}, train_acc{train_acc:.4f}, val_acc{val_acc:.4f}) if val_acc best_acc: best_acc val_acc torch.save(model.state_dict(), best_model.pth)我只跑了20个epoch做基线验证。CIFAR-10这个任务上ViT-Base微调20轮之后验证精度大概能到90%以上如果你把训练策略调得更精细还能更高。5.4 训练中观察到的现象我在实际跑的过程中有几个体会前两三个epoch准确率会快速上升这是预训练权重在下游任务上“热身”的过程。如果loss一直不降先检查学习率是不是太大或者预处理中Normalize的均值方差是不是写错了。验证集准确率偶尔会抖动不用慌只要整体趋势向上就行最终模型以验证集最优为准。6. 可视化输出训练曲线、混淆矩阵与注意力热力图模型训练完光有一个准确率数字远远不够。你需要把训练过程、分类结果、模型关注区域都画出来。这个章节也是“从环境搭建到可视化输出”的最后一环。6.1 绘制训练曲线训练曲线能直观反映收敛状态。最简单的方式是用matplotlib记录每个epoch的train_loss和val_accimport matplotlib.pyplot as plt # history {train_loss: [], val_acc: []} plt.figure(figsize(12, 4)) plt.subplot(1, 2, 1) plt.plot(history[train_loss], labeltrain_loss) plt.xlabel(epoch) plt.ylabel(loss) plt.legend() plt.subplot(1, 2, 2) plt.plot(history[val_acc], labelval_acc) plt.xlabel(epoch) plt.ylabel(accuracy) plt.legend() plt.tight_layout() plt.savefig(training_curves.png)如果你习惯用TensorBoard也可以在训练循环里加入SummaryWriter记录。TensorBoard的优势是能实时看曲线不用等训练结束。6.2 绘制混淆矩阵混淆矩阵能告诉你模型到底把哪些类搞混了。我用scikit-learn计算再用seaborn画图。import numpy as np from sklearn.metrics import confusion_matrix import seaborn as sns all_preds, all_labels [], [] model.eval() with torch.no_grad(): for x, y in val_loader: x x.cuda() logits model(x) preds logits.argmax(dim1).cpu().numpy() all_preds.extend(preds) all_labels.extend(y.numpy()) cm confusion_matrix(all_labels, all_preds) plt.figure(figsize(8, 8)) sns.heatmap(cm, annotTrue, fmtd, cmapBlues) plt.xlabel(Predicted) plt.ylabel(True) plt.savefig(confusion_matrix.png)如果发现某两个类互相混淆严重通常说明它们在视觉上确实相似或者样本数量不均。这时候可以针对性增加数据增强或者补充这类图片的训练量。6.3 注意力热力图模型到底在看哪里这是ViT最吸引人的地方。因为Transformer的注意力权重天然可以提取出来告诉你模型做判断时重点参考了哪些patch。你可以用下面的思路获取最后一层Encoder的注意力图import torch.nn.functional as F # 注册一个钩子获取最后一个Transformer Block的注意力权重 attention_maps [] def hook_fn(module, input, output): # output 里包含 attention 权重需要按实际模型结构取 attention_maps.append(output) # 以timm的ViT为例通常可以这样访问 model.blocks[-1].attn.register_forward_hook(hook_fn)更务实的做法是直接使用timm内部提供的get_intermediate_layers方法它能返回每个block的输出特征。取最后一层所有注意力头的平均值按patch的14x14网格还原成空间尺寸再放大到原图大小。with torch.no_grad(): # 输入一张(1,3,224,224)的图 feats model.get_intermediate_layers(x, n1)[0] # [1, 197, 768] # 得到的是最后一层的特征不是注意力权重如果你只是想快速看到模型关注区域更实用的方案是使用GradCAM类方法作用于ViT的最后一个卷积层或注意力块。对于ViT可以基于最后一层所有Token对CLS Token的注意力权重来生成热力图。基本思路是取出最后一层注意力权重中CLS Token对应的那一行丢掉CLS自己剩下的196个值reshape成14x14上采样回224x224再与原图叠加。这个热力图能非常直观地告诉你模型是否关注了目标区域。我在实际项目里发现模型分类正确时热力图通常集中在主体物体上如果热力图乱飘那基本可以断定模型在“抄近道”靠背景信息作弊这时候就要考虑是否数据有泄漏或者增强方式是否需要调整。7. 我踩过的几个坑ViT训练中反直觉的细节最后这一节不写教程式的标准步骤纯粹分享我实际踩过、也花了不少时间排查的坑。每一条都应该能帮你在自己的项目中少走弯路。7.1 预训练模型的输入尺寸必须对齐第一次用ViT时我直接拿256x256的图像输入心想“分辨率大一点信息多一点总没坏处”。结果模型直接报维度错误。原因在于位置编码的数量是训练时定死的224分辨率对应196个patch256分辨率则对应256个patch位置编码矩阵对不上。timm里如果你传pretrainedTrue它会用224对应的编码。想用大图就得用支持可变分辨率的位置编码插值方案或者干脆选针对384分辨率预训练的变体。我后来的习惯是除非特别需要否则就用224x224省心又够用。7.2 学习率给大了一轮就崩卷积网络微调时我用1e-4经常没事换到ViT之后第一次也照抄了这个值结果训练到第5轮loss开始往上飙。后来检查发现ViT对学习率的容忍度明显更低。把学习率降到3e-5之后训练一下就稳了。我的经验值是ViT微调学习率从1e-5到5e-5之间起步比较安全。你要是实在想冲更快也建议先用3e-5跑通一轮确认loss下降稳定再加。7.3 忘了切换训练/验证模式这个错误特别低幼但真的容易忽略。验证阶段如果不调用model.eval()BatchNorm和Dropout的行为会和训练时不一样准确率可能忽高忽低。ViT里虽然BatchNorm用得少但Dropout依然存在影响不小。7.4 显存溢出不一定是batch size的锅有一次我训练中途报显存溢出第一反应是减小batch size结果发现是验证阶段忘记写torch.no_grad()导致推理时也保留了完整计算图。这个坑很隐蔽因为训练阶段显存本来就用得多溢出时不容易想到是验证代码的问题。7.5 自己数据集类别数和预训练头不一致别忽略很多人觉得timm里num_classes参数传一下就行其实替换分类头之后新初始化的权重是随机值前几轮梯度主要集中在这一层。如果你的数据集很小新分类头的训练可能不够充分。我的解决办法是给新分类头单独设一个稍大的学习率或者先把分类头单独训练几个epoch再全量微调。这个小技巧在数据量少的时候往往会带来意外惊喜。8. 从基线到进阶后续还能怎么扩展到这里你已经跑通了一条完整的ViT图像分类链路也拿到了可视化结果。但我还想多说几句后续的扩展方向因为ViT的想象力远不止“替掉ResNet做分类”这么简单。8.1 把注意力热力图做成线上诊断工具我在真实项目中会把注意力热力图写入TensorBoard每轮训练结束随机抽几张验证图看模型的注意力是否稳定落在目标上。这个操作成本极低但对模型调优帮助很大。你会发现有些时候模型“准确率高”其实是一种假象它靠的是背景或水印。这类问题只有可视化才能看出来。8.2 尝试更强的数据增强策略如果你想在CIFAR-10或类似小数据集上冲击更高精度可以在预处理里加入RandomResizedCrop、RandAugment、MixUp和CutMix。我自己实验下来这些增强策略对ViT的帮助比传统卷积网络更明显。原因是ViT的归纳偏置少需要靠数据增强“补课”的地方更多。8.3 从分类到其他任务ViT的分类头拿掉之后它就是一个通用的视觉特征提取器。你可以在此基础上做图像检索、零样本分类甚至接入文本编码器做图文匹配。如果你之后要接触多模态模型ViT的架构知识几乎是必需品。9. 一些个人经验与最后的建议折腾ViT的这段时间我最深的体会是ViT本身不复杂复杂的是你脑子里的那套“卷积思维”要切换过来。刚开始我老是下意识地问“感受野多大”“池化在哪”但这些概念在ViT里都被重新定义了。放下这些旧框架之后代码写起来反而更简单因为整个流程就是“切块、编码、Transformer、接分类头”。如果你准备动手实践我的建议是别急于一次跑通所有内容。按照最简单的路径来先加载预训练模型在CIFAR-10上跑通一个epoch确认训练曲线在下降接着加上完整的验证和模型保存最后再尝试可视化。一步一步来出了问题时定位会非常清晰。最后分享一个很实用的小技巧加载完预训练权重后别急着改任何训练参数先用一张图片做一次前向推理直接看图能不能输出合理的类别概率。如果这一步就错了说明环境或者模型加载有问题后面怎么调都白搭。我的所有失败项目几乎都在这类“前置检查没做”的情况下浪费过时间。希望这篇文章能让你少走一些弯路。如果跑的过程中遇到什么奇怪的报错欢迎回来对照着第七章的坑查一遍大概率问题就出在那里。