简介基于Vision TransformerViT的CIFAR-10图像分类训练与验证Python源码面向深度学习初学者、计算机视觉相关专业学生以及正在准备课程设计或毕业设计的人员旨在帮助读者快速掌握ViT在小型图像分类任务中的完整实践流程。压缩包共2个文件以Python脚本为主负责ViT模型定义、数据加载及训练验证逻辑另附structure.txt结构说明清晰梳理目录与模块关系。包体仅2KB轻量精简适合直接查看核心代码与工程组织方式。目前已有526人学习下载代码经实际运行验证可稳定完成数据加载、模型训练与验证评估等环节。源码覆盖从图像预处理、patch嵌入、Transformer编码器到分类头的完整实现路径既可作学习样例也可作为基线对比CNN方案借助txt中的说明可快速定位关键模块便于在此基础上修改或扩展实验。1. ViT在CIFAR10上到底行不行小模型、小数据的实际表现如果你拿ViTVision Transformer直接套CIFAR10会撞上一个很尴尬的事实同样的算力下ResNet18稍微调调就到92%而ViT的复现脚本往往卡在90%附近甚至更低于是不少人扭头就把ViT丢进“在小数据集上不靠谱”的垃圾桶。但从我实际跑过的源码来看问题几乎不在注意力机制本身而在位置编码怎么初始化、学习率要不要warmup、dropout和weight decay怎么配。CIFAR10的图只有32x32用patch_size4切成64个tokenViT照样能涨到93%以上。这篇就是把我自己整理的一套“基于Vit实现CIFAR10分类数据集的训练和验证python源码”拆给你看从依赖安装到模型实现再到训练验证循环和5个我踩过的坑全部是能直接复制跑的代码。2. 搭建ViT跑CIFAR10的最小环境从依赖安装到DataLoader配置2.1 python环境与依赖python 3.8、PyTorch和torchvision的搭配先把环境立住。常见的做法是直接用conda建一个干净环境python版本选3.8或者3.10PyTorch选2.xtorchvision跟着PyTorch走。我这里给你一个能直接用的安装命令conda create -n vit_cifar python3.8 -y conda activate vit_cifar pip install torch2.1.2 torchvision0.16.2 --index-url https://download.pytorch.org/whl/cu118 pip install numpy matplotlib tqdm tensorboard说明一下torch和torchvision的版本要对应我是用--index-url指定了CUDA 11.8的源如果你机器上没有NVIDIA显卡把最后一个index-url去掉装CPU版也能跑只是慢一些。CIFAR10训练本身不重CPU跑几十个epoch也不是不能忍但 ViT 的训练实际上很吃矩阵运算有卡还是用卡。numpy是给数据预处理和标签操作用的matplotlib和tensorboard留着后面做可视化验证。这里要提一下很多初学者在python环境上的翻车点不要直接pip install torch拉到最新版新版torchvision可能要求更高版本的python你如果是python 3.8拉到不兼容的包会当场报错。装完之后用python -c import torch; print(torch.__version__)测一下能输出版本号再往下走。2.2 用torchvision按官方方式加载CIFAR10数据集下载与类别标签核对环境就绪之后最稳的数据加载姿势是走torchvision的datasets.CIFAR10接口它负责下载、解压、按训练集验证集切分。下面这段代码就是完整的pytorch加载cifar10的方式import torch from torchvision import datasets, transforms from torch.utils.data import DataLoader transform_train transforms.Compose([ transforms.RandomCrop(32, padding4), transforms.RandomHorizontalFlip(), transforms.ToTensor(), transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2470, 0.2435, 0.2616)), ]) transform_val transforms.Compose([ transforms.ToTensor(), transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2470, 0.2435, 0.2616)), ]) train_dataset datasets.CIFAR10( root./data, trainTrue, downloadTrue, transformtransform_train ) val_dataset datasets.CIFAR10( root./data, trainFalse, downloadTrue, transformtransform_val ) print(train_dataset.classes) # [airplane, automobile, bird, cat, deer, dog, frog, horse, ship, truck]这段的逻辑很简单训练集做RandomCrop和RandomHorizontalFlip做数据增强验证集只做ToTensor和Normalize不做任何增强保证验证指标的稳定性。Normalize的均值和方差是CIFAR10官方统计好的不要自己算自己算也能用但没必要。注意root./data第一次跑会自动下载到当前目录的data文件夹里大概160MB。如果你在公司内网或者网络受限下载卡住了可以去网上找CIFAR10的压缩包手动放到./data目录下再跑这个接口会识别已经存在的文件不会重复下载。标签顺序就是上面那10个类这个顺序经常被忽略后面做混淆矩阵时就会用上建议先打印出来看一眼。2.3 DataLoader的关键参数num_workers、pin_memory与验证集不要drop_last数据加载这块直接照抄这个配置train_loader DataLoader( train_dataset, batch_size128, shuffleTrue, num_workers4, pin_memoryTrue, drop_lastTrue ) val_loader DataLoader( val_dataset, batch_size256, shuffleFalse, num_workers4, pin_memoryTrue, drop_lastFalse )参数说明batch_size128是CIFAR10上比较稳的选择ViT-small级别的模型显存占用不高128够吃shuffleTrue只给训练集验证集不需要打乱保持固定顺序方便复现指标drop_lastTrue只给训练集防止最后一个batch样本数不足导致BatchNorm或LayerNorm统计量抖动验证集必须drop_lastFalse因为验证集总共10000张256整除不了剩下几张小batch也要参与计算不能丢。num_workers4在Linux上没问题但在Windows上经常因为多进程数据加载报错如果报BrokenPipe或者内存持续飙升把它改成2就老实了。pin_memoryTrue配合GPU训练时能把数据从CPU拷贝到GPU的速度提一点显存不紧张就开着。3. 手写Vision Transformer关键模块Patch Embedding到Transformer Encoder3.1 Patch Embedding用Conv2d一步实现为什么卷积层能当线性投影用ViT的思想是把图像切成一堆patch再把每个patch线性投影成embedding。最常见的实现不是真的去切图而是用一个kernel_sizestridepatch_size的Conv2d一步搞定。下面就是把patch embedding和位置编码封在一起的源码import torch import torch.nn as nn class PatchEmbed(nn.Module): def __init__(self, in_channels3, image_size32, patch_size4, embed_dim192): super().__init__() self.image_size image_size self.patch_size patch_size self.grid (image_size // patch_size) ** 2 # CIFAR10: 8*864个patch self.proj nn.Conv2d( in_channels, embed_dim, kernel_sizepatch_size, stridepatch_size ) def forward(self, x): # x: [B, 3, 32, 32] - [B, embed_dim, 8, 8] x self.proj(x) # 展平成token序列: [B, embed_dim, 8, 8] - [B, 64, embed_dim] x x.flatten(2).transpose(1, 2) return x这段代码的逻辑Conv2d的padding为0kernel_size和stride都等于patch_size输入32x32的图经过卷积后空间尺寸变成8x8同时通道数从3变成embed_dim相当于每个4x4的patch被线性投影成了一个192维的向量。然后flatten(2)把8x8展平成64个位置transpose(1, 2)把维度调成[B, 64, embed_dim]这就是标准的token序列。参数说明patch_size4是关键CIFAR10的图太小如果照搬ImageNet常用的patch_size16整张图只剩2x24个tokenTransformer基本没法做注意力。embed_dim192是偏小的维度对应ViT-Small缩减版。如果你想把模型加大可以改到256或384但训练时间和显存都会涨。3.2 位置编码可学习位置编码与图像大小变化时的处理Transformer本身没有顺序概念patch的顺序信息全靠位置编码。ViT论文里用的是可学习位置编码这也是一种“黑匣子”你不需要手工设计初始化成0或者用正态分布随机初始化训练时会自己学出来。class PositionEmbedding(nn.Module): def __init__(self, num_patches64, embed_dim192): super().__init__() # 可学习位置编码形状是 [1, num_patches, embed_dim] self.pos_embed nn.Parameter( torch.zeros(1, num_patches, embed_dim) ) nn.init.trunc_normal_(self.pos_embed, std0.02) def forward(self, x): # x: [B, 64, 192] return x self.pos_embed这里有两个容易出问题的点。第一num_patches必须和PatchEmbed产出的token数一致CIFAR10就是64别写成196那是ImageNet的尺寸。第二如果以后你想把图像分辨率提高比如从32改成64patch数量就会从64变成256此时可学习位置编码的形状对不上常见的做法是对位置编码做双线性插值或者干脆重新训练。对于CIFAR10这种固定尺寸的任务直接用可学习编码就行不要去折腾插值。初始化用trunc_normal_是ViT源码里的习惯std取0.02比默认的均匀分布收敛更稳这个数值我实测过影响不算大但值得保留。3.3 Transformer Encoder与分类头LayerNorm/Dropout放在哪里Encoder部分就是把标准Transformer的Encoder层堆起来。ViT和NLP里的BERT有个细节差异ViT在patch embedding之后不加[CLS]token而是用全局平均池化Mean Pooling把64个token压成一个向量再过分类头这个选择对CIFAR10这种小图反而更稳。class TransformerBlock(nn.Module): def __init__(self, embed_dim192, num_heads6, mlp_ratio4.0, dropout0.1): super().__init__() self.norm1 nn.LayerNorm(embed_dim) self.attn nn.MultiheadAttention( embed_dim, num_heads, dropoutdropout, batch_firstTrue ) self.norm2 nn.LayerNorm(embed_dim) mlp_hidden int(embed_dim * mlp_ratio) self.mlp nn.Sequential( nn.Linear(embed_dim, mlp_hidden), nn.GELU(), nn.Dropout(dropout), nn.Linear(mlp_hidden, embed_dim), nn.Dropout(dropout) ) def forward(self, x): # Pre-LN结构: 先norm再进注意力训练更稳定 x x self.attn(self.norm1(x), self.norm1(x), self.norm1(x))[0] x x self.mlp(self.norm2(x)) return x class ViTClassifier(nn.Module): def __init__(self, in_channels3, image_size32, patch_size4, embed_dim192, depth6, num_heads6, num_classes10, dropout0.1): super().__init__() self.patch_embed PatchEmbed(in_channels, image_size, patch_size, embed_dim) num_patches (image_size // patch_size) ** 2 self.pos_embed PositionEmbedding(num_patches, embed_dim) self.blocks nn.Sequential(*[ TransformerBlock(embed_dim, num_heads, dropoutdropout) for _ in range(depth) ]) self.norm nn.LayerNorm(embed_dim) self.head nn.Linear(embed_dim, num_classes) self.dropout nn.Dropout(dropout) def forward(self, x): x self.patch_embed(x) x self.pos_embed(x) x self.dropout(x) x self.blocks(x) x x.mean(dim1) # 全局平均池化取所有token的均值 x self.norm(x) x self.head(x) return x逻辑说明每个TransformerBlock用的都是Pre-LN结构先LayerNorm再做注意力这个和Post-LN相比梯度传播更顺训练大深度时不容易爆炸。nn.MultiheadAttention里batch_firstTrue是PyTorch 2.x的写法这样输入输出都是[B, seq_len, dim]省去了维度交换的麻烦。注意力之后接残差连接MLP块里先升维再降维激活函数用GELU而不是ReLU这是Transformer家族的标准配置。分类头那里x.mean(dim1)就是全局平均池化它对所有patch token取平均。3.4 ViT-CIFAR10超参表这些参数决定了收敛速度把你的模型实例化参数照这个表走参数取值说明image_size32CIFAR10原始分辨率patch_size44x4像素一个token共64个tokenembed_dim192token嵌入维度加大到256/384能提点但慢depth6Transformer Encoder层数4~8之间调num_heads6注意力头数embed_dim必须能整除mlp_ratio4.0MLP隐藏层维度embed_dim*4dropout0.1位置编码后和MLP内的dropout概率num_classes10CIFAR10固定embed_dim192, num_heads6是一个经验搭配192除以6等于32每个head分到32维这是注意力头比较舒服的宽度。如果你把embed_dim改到256num_heads就建议改成8。depth取6层对CIFAR10来说已经能装下足够的信息再往上加深收益会被过拟合吃掉训练时间却翻倍。dropout在CIFAR10这种小数据上不能省0.1是起步值如果你发现训练集准确率明显高于验证集试着把它加到0.2。4. 训练与验证循环损失函数、学习率调度与checkpoint保存4.1 训练一个epoch的完整骨架AMP、梯度裁剪与loss计算模型写好了下面就是训练循环。这里我直接给出简洁、又能复现的写法用自动混合精度加速训练加梯度裁剪防止loss跳出悬崖import torch import torch.nn as nn from tqdm import tqdm def train_one_epoch(model, loader, optimizer, criterion, scaler, device): model.train() total_loss 0.0 correct 0 total 0 for images, labels in tqdm(loader, desctraining): images, labels images.to(device), labels.to(device) optimizer.zero_grad() with torch.autocast(device_typecuda, dtypetorch.float16): logits model(images) loss criterion(logits, labels) scaler.scale(loss).backward() scaler.unscale_(optimizer) nn.utils.clip_grad_norm_(model.parameters(), max_norm1.0) scaler.step(optimizer) scaler.update() total_loss loss.item() * images.size(0) _, preds logits.max(dim1) correct preds.eq(labels).sum().item() total labels.size(0) return total_loss / total, correct / total代码说明torch.autocast和scaler是PyTorch原生AMP的搭配混合精度能把训练速度提升30%~50%在ViT这种计算密集型模型上特别明显。如果你用的是CPU只有device_typecuda会报错把AMP整段换成普通的前向反向就行。梯度裁剪max_norm1.0是有必要的ViT在训练初期偶尔会出现loss突增裁剪之后模型不会因为一个异常梯度直接崩掉。criterion建议用nn.CrossEntropyLoss()这是分类问题的标准选择。如果你的数据增强用了MixUp后面避坑章节会提标签就需要改成软标签此时nn.CrossEntropyLoss()不能直接用要换成KL散度版本我在第5章专门说。4.2 学习率调度策略warmup cosine decayViT训练和CNN最大的区别就是对学习率敏感。业界最稳的方案是先做一个短warmup线性升温再用cosine调度缓慢下降。我把调度器和训练循环的骨架给你import math from torch.optim import AdamW from torch.optim.lr_scheduler import LambdaLR def lr_lambda(step, warmup_steps, total_steps): if step warmup_steps: return (step 1) / warmup_steps progress (step - warmup_steps) / max(1, total_steps - warmup_steps) return 0.5 * (1.0 math.cos(math.pi * progress)) model ViTClassifier().to(device) optimizer AdamW(model.parameters(), lr1e-3, weight_decay0.05) total_steps len(train_loader) * epochs warmup_steps int(total_steps * 0.05) scheduler LambdaLR(optimizer, lr_lambdalambda step: lr_lambda(step, warmup_steps, total_steps))参数说明lr1e-3搭配weight_decay0.05是ViT原论文的经典组合相比CNN常用的SGD0.9动量AdamW在Transformer上收敛更顺。warmup_steps取总迭代步数的5%比如你跑100个epochwarmup大约占5个epoch学习率从0线性涨到1e-3然后再经过95个epoch的cosine下降回到接近0。这里强调一个区别LambdaLR的step是从0开始的迭代步数不是epoch数训练循环里每轮迭代结束之后记得调scheduler.step()这里很容易漏。4.3 验证循环用确定性模式评估模型验证代码没有太多玄学但必须注意模型模式切换torch.no_grad() def validate(model, loader, criterion, device): model.eval() total_loss 0.0 correct 0 total 0 for images, labels in tqdm(loader, descvalidating): images, labels images.to(device), labels.to(device) with torch.autocast(device_typecuda, dtypetorch.float16): logits model(images) loss criterion(logits, labels) total_loss loss.item() * images.size(0) _, preds logits.max(dim1) correct preds.eq(labels).sum().item() total labels.size(0) return total_loss / total, correct / totalmodel.eval()会关掉Dropout和LayerNorm的统计更新这一步不做验证准确率会忽高忽低很多人以为模型训练出问题了其实是这个细节。torch.no_grad()显式关闭梯度图验证阶段的显存占用会大幅下降能防止验证时OOM。另外验证的时候不需要梯度裁剪和AMP的scaler更新按照上面的代码只做前向和统计就行。每轮epoch结束后把这个函数的返回值打印出来观察验证loss和准确率。4.4 三个必调参数batch size、weight decay、dropout训练循环跑起来之后你会遇到一堆调节项我只讲三个影响最大的第一个是batch size。CIFAR10上128是一个甜点值显存占用不大每个step的梯度噪音适中。如果你把它降到32训练会变得很不稳定ViT这种模型对小batch很敏感升到256以上收敛变慢需要同步调大学习率否则训练速度看着快、实际能达到的准确率却下滑。第二个是weight decay。AdamW的0.05是ViT原论文的值正则强度不小换成0.01会稍微提高训练集准确率但验证集可能降0.3%左右如果你想在CIFAR10上追求极限验证准确率0.05不用动。第三个是dropout。如果验证准确率明显低于训练集典型过拟合先加dropout到0.2或0.3比改模型结构更快见效。注意位置编码后的self.dropout是token维度的它的意义是给整个序列注入噪声MLP里的Dropout是在特征维度上做两者作用范围不同调节时先动后者。5. 避坑与排查5个实测踩坑记录5.1 训练loss不下降位置编码和token维度匹配出错了现象loss从2.3附近掉到2.0之后就不再动准确率永远停在10%附近相当于随机猜。原因最常见的是两个。一是位置编码的形状和patch数量不一致PyTorch在相加时不报错但会走广播机制把位置编码错位加到token上模型学不到位置信息二是分类头之前用了全局平均池化但代码写成了x.mean(dim0)相当于对batch维度做了平均整个backbone在训练时会看到一团浆糊。解决打印模型每一步输出的shapex.mean()只在最后一维上操作。从PatchEmbed开始分别打印x.shape确认是[B, 64, 192]再到[B, 192]。位置编码那边初始化后做一个单步前向确保没有报错且输出的第一个token确实带上了位置信息。5.2 验证集准确率来回震荡模型推理时开着dropout现象训练集准确率平滑上升验证集准确率一会上90%一会掉到70%像心跳一样波动。原因这就是我在4.3里强调过的model.eval()没写或者写了但没调用对位置。Dropout在eval模式下虽然被关掉但如果你把validate函数里的model.eval()注释掉了模型每次前向都用不同的dropout mask验证结果自然震荡。解决validate函数进循环之前确保调用model.eval()训练之前确保调用model.train()。还有一个更隐蔽的坑如果你把模型包在nn.DataParallel或者torch.compile里eval/train模式的切换会传染给子模块此时用model.module.eval()或model.eval()统一处理不要混着写。5.3 Windows下DataLoader多进程报错把num_workers改成0或2现象程序一启动训练立刻弹BrokenPipe或者RuntimeError: DataLoader worker (pid(s)) exited unexpectedly网上查半天找不到原因。原因Windows的多进程启动方式和Linux不同num_workers4在Windows上经常触发子进程的pickle坑特别是数据加载函数里用了lambda或局部类时子进程无法序列化这些对象。解决既然只是CIFAR10这种小数据集数据加载本身不是瓶颈直接num_workers2或者干脆改成num_workers0让主进程同步加载。0会慢一些但最省心。另外把训练脚本放到if __name__ __main__:块里这是Windows多进程的硬性要求。5.4 MixUp剪裁后标签不匹配概率标签与验证指标的冲突现象你给训练数据加了MixUp增强训练loss下降得很漂亮但每次validate算出来的准确率都只有70%多怎么调都上不去。原因MixUp会把两张图的标签按比例融合成软标签比如0.7*猫 0.3*狗。但很多复现脚本只改了图片生成逻辑没改loss计算方式还在用nn.CrossEntropyLoss()硬套软标签或者验证时直接用argmax去比对软标签这会在指标上制造“假翻车”——模型其实已经学得不错只是你用错了评估方式。解决训练阶段用KL散度做软标签loss验证阶段用原始硬标签算准确率。也就是说训练集加载器里做MixUp并保留融合权重验证集加载器不做MixUp保持原始标签。有了这个前提再去看准确率才可信。5.5 训练中loss突然变NaN学习率过于激进现象跑到第十几个epochloss突然跳到nan之后一直nan准确率归零。原因AdamW的学习率1e-3本身没问题问题是warmup没做或太短。0学习率瞬间跳到1e-3模型浅层参数的更新步长过大梯度碰到数值溢出另外AMP混合精度下fp16的梯度下溢也可能触发nan。解决把warmup_steps的比例从5%提到10%loss还是nan就降学习率到5e-4。梯度裁剪保留着它防的是“梯度爆炸”对“梯度溢出”也有一定兜底。如果是在AMP开启时出现nan可以试scaler.set_growth_factor(2.0)调低梯度scaler的增长倍数这招能解决一部分fp16的数值问题。6. 验证ViT真实水平的两个可视化技巧混淆矩阵与学习率曲线6.1 混淆矩阵看模型在哪几个类别上互相混淆准确率只是一个数字真正想定位模型的弱点还是要看混淆矩阵。CIFAR10里最容易搞混的是猫和狗、汽车和卡车如果混淆矩阵对角线意外地不干净说明模型学到的是纹理特征而不是语义特征。下面是生成混淆矩阵的代码建议每训练完一轮就画一次看对角线变化import numpy as np import matplotlib.pyplot as plt from sklearn.metrics import confusion_matrix, ConfusionMatrixDisplay def plot_confusion(model, val_loader, device): model.eval() all_preds, all_labels [], [] with torch.no_grad(): for images, labels in val_loader: images images.to(device) logits model(images) preds logits.argmax(dim1).cpu().numpy() all_preds.extend(preds) all_labels.extend(labels.numpy()) cm confusion_matrix(all_labels, all_preds) disp ConfusionMatrixDisplay(confusion_matrixcm, display_labelstrain_dataset.classes) disp.plot(xticks_rotation45) plt.tight_layout() plt.savefig(confusion_matrix.png, dpi150)这段代码比较直观把验证集所有预测结果收集起来用sklearn画矩阵。注意train_dataset.classes要传进去作为标签名不能用数字0~9当显示标签否则你大概率看不懂错在哪。比如deer和horse这两种都有四条腿和棕色纹理如果它们互相混淆的比例很高就可以考虑在数据增强里加点颜色抖动帮助模型把颜色和形状解耦。6.2 学习率曲线与训练记录把“玄学”变成可排查的日志最后一招也是我自己的一个习惯每个epoch把train_loss、train_acc、val_loss、val_acc、learning_rate这五个值记录成csv训练结束后画成曲线。这个习惯帮我发现了不少问题比如val_loss在第50个epoch开始回升而train_loss还在降那就是过拟合信号加weight decay或dropout。import csv log_path training_log.csv with open(log_path, w, newline) as f: writer csv.writer(f) writer.writerow([epoch, train_loss, train_acc, val_loss, val_acc, lr]) # 每个epoch结束后追加一行 # writer.writerow([epoch, train_loss, train_acc, val_loss, val_acc, cur_lr])学习率曲线重点看warmup结束后的第一个拐点如果那附近验证准确率跟着明显跳升说明模型前期卡在局部平缓区warmup起了作用如果拐点后验证准确率反而下跌说明初始学习率还是偏高下次调低一半。这套源码跑通之后的正常结果是100个epoch左右验证准确率在92%~93%之间。ViT在小数据集上不会突破94%太多如果你的目标只是分类精度ResNet系列依然是性价比之王但如果你想研究注意力机制、可视化attention map或者想把分类头换成CLIP式的图文对齐这套CIFAR10上的ViT训练验证源码就是最轻量、最能自由改造的起点。我自己的经验是把depth调到8、embed_dim保持192时在CIFAR10上能到93.5%但训练时间涨了差不多40%。你在跑的时候记住验证集指标波动超过0.5%就优先检查eval模式是否关闭了dropout训练不收敛就先把warmup拉长这些坑我都替你踩过了希望帮到你。本文还有配套的精品资源点击获取