简介基于视觉TransformerViT实现CIFAR-10分类数据集训练与验证的Python源码以CIFAR-10为基准图像分类任务面向计算机、人工智能、通信工程、自动化等专业的在校学生、教师和企业开发者也适合有一定基础的学习者作为进阶参考或快速开展图像分类实验。压缩包体积仅2KB包含两个文件一个完整的模型训练与验证Python脚本一个项目结构说明文本文件方便快速定位关键代码与模块关系。代码经过完整测试且运行稳定作者将其用于毕业设计并取得96分的评审成绩因此可直接复现CIFAR-10图像分类流程也可作为课程设计、毕业设计、项目初期演示的起点。目前已有525人学习使用读者可在此基础上修改数据预处理、网络结构或训练参数进一步拓展到其他分类任务是一个轻量、易上手的ViT实践资源。1. 基于Vit的CIFAR10分类这份源码到底教会你什么我拆过很多分类项目的源码但每次看到有人拿 ViT 直接怼 CIFAR10 还是会多看两眼。原因很简单ViT 是给大图224×224 以上设计的CIFAR10 是 32×32 的小图直接套用原版 ViT 效果往往不如 ResNet这不是模型差是输入处理方式没跟上。这份vit_cifar10-master.zip里正好解决的就是这个问题——它用一套可调的 ViT 结构完成 CIFAR10 从训练到验证的完整闭环核心文件就一个Vit.py加上配套训练脚本。适合两类人一类是做计科、AI 相关毕设或课程设计的学生需要一份能跑通、能改、能讲清楚原理的基线代码另一类是刚接触 Transformer 做视觉任务、想弄明白 patch 是怎么变成 token 的入门者。这份资源不复杂但把 ViT 最关键的数据流、损失计算、训练循环都摊开在你面前了。2. Vit.py 的四个关键结构从 Patch Embedding 到分类头2.1 为什么 ViT 能处理 32×32 的小图Patch 划分是前提ViT 的核心假设是“图像可以像句子一样被切块”。对于 224×224 的 ImageNet 图像原版用 16×16 的 patch得到 196 个 token。但 CIFAR10 只有 32×32如果用 16×16 patch就只剩 4 个 tokenTransformer 的注意力在这种序列长度下根本学不动。所以Vit.py里第一个可调参数就是patch_size。常见做法是把 patch 设成 4×4这样 32×32 的图像被切成 64 个 patch序列长度虽然比 ImageNet 短很多但足够注意力机制发挥。这份源码里默认的image_size32, patch_size4就是为 CIFAR10 定制的。如果你想迁移到自己的小图数据集这一步是最先要确认的——如果图本身小于 patch_size代码会直接报维度错误如果图和 patch_size 不匹配后面的 Reshape 一定会翻车。来看Vit.py里最核心的 Patch Embedding 实现class PatchEmbed(nn.Module): def __init__(self, image_size32, patch_size4, in_channels3, embed_dim256): super().__init__() self.image_size image_size self.patch_size patch_size self.patch_num (image_size // patch_size) ** 2 # 8*864 self.proj nn.Conv2d(in_channels, embed_dim, kernel_sizepatch_size, stridepatch_size) def forward(self, x): # x: [B, 3, 32, 32] - [B, 64, 256] x self.proj(x) # [B, 256, 8, 8] x x.flatten(2) # [B, 256, 64] x x.transpose(1, 2) # [B, 64, 256] return x这里没有用传统的“切块再展平再线性映射”而是直接用Conv2d来完成。卷积核大小和步长都等于patch_size等价于把每个 patch 做一次线性变换输出维度就是embed_dim。这一步有两点值得注意patch_num由图像尺寸和 patch 尺寸共同决定改成别的分辨率时要同步算一遍。flatten(2)是从宽高维度展平transpose(1,2)是把通道维提到前面顺序反了会导致 token 的特征错位训练出来准确率可能只有 30%但不会报错是最难排查的隐性 bug。2.2 Class Token 和 Position EmbeddingViT 的“全局池化”替代方案CNN 分类通常用全局平均池化把特征图变成向量ViT 用的是一个更“暴力”的办法在序列最前面拼一个可学习的cls_token让这个 token 自己学会聚合整张图的信息。最后分类头只取这个 token 对应的输出其余 token 全部丢掉。Vit.py里这部分代码长这样class ViT(nn.Module): def __init__(self, ...): self.cls_token nn.Parameter(torch.zeros(1, 1, embed_dim)) self.pos_embed nn.Parameter(torch.zeros(1, 1 self.patch_embed.patch_num, embed_dim)) nn.init.trunc_normal_(self.cls_token, std0.02) nn.init.trunc_normal_(self.pos_embed, std0.02) def forward(self, x): B x.shape[0] x self.patch_embed(x) # [B, 64, 256] cls_tokens self.cls_token.expand(B, -1, -1) # [B, 1, 256] x torch.cat([cls_tokens, x], dim1) # [B, 65, 256] x x self.pos_embed # [B, 65, 256] ... return x[:, 0] # 取 cls_tokenpos_embed的形状是(1, 1patch_num, embed_dim)1 是 cls_token 占的位置后面 64 是图像 patch 的位置。这里初始化用的是trunc_normal_标准差 0.02这是从 DeiT 和 MAE 里继承下来的习惯。值得注意如果图像尺寸不是 32×32patch_num变了pos_embed的第二个维度也要跟着变否则加法直接报 shape mismatch。有些简化版实现会在x x self.pos_embed之前对pos_embed做repeat(B,1,1)这个源码里直接用广播因为pos_embed第一个维度是 1PyTorch 会自动扩展不用手动复制。2.3 Transformer Encoder自注意力层的维度必须牢记ViT 的 Encoder 就是标准 Transformer 的编码器由多层MSA MLP堆叠。Vit.py里通常会把depth设为 68num_heads设为 8mlp_ratio设为 4。这组参数是 CIFAR10 上效果和速度比较平衡的搭配。看一段典型的 Encoder 层实现class TransformerEncoder(nn.Module): def __init__(self, embed_dim256, num_heads8, depth6, mlp_ratio4.0, dropout0.1): super().__init__() self.layers nn.ModuleList([ nn.TransformerEncoderLayer( d_modelembed_dim, nheadnum_heads, dim_feedforwardint(embed_dim * mlp_ratio), dropoutdropout, batch_firstTrue, activationgelu ) for _ in range(depth) ]) self.norm nn.LayerNorm(embed_dim) def forward(self, x): for layer in self.layers: x layer(x) return x这里直接用了 PyTorch 自带的nn.TransformerEncoderLayer省去了手写 QKV 矩阵的麻烦。注意batch_firstTrue否则输入输出维度是(seq, batch, embed)跟前面 Patch Embedding 输出的(B, seq, embed)对不上。dim_feedforward是 MLP 隐藏层维度一般取embed_dim * 4太小会降低模型容量太大在小数据集上容易过拟合。还有一个易错点nn.TransformerEncoderLayer默认的dropout只作用在残差块的输出和 MLP 内部如果你要加一个整体的dropout记得放在 Encoder 之后、分类头之前而不是加在每层里。2.4 分类头和损失计算10 类输出的最后一公里分类头就是一层线性层把embed_dim映射到 10。训练时用交叉熵损失验证时取argmax。class ViTClassifier(nn.Module): def __init__(self, embed_dim256, num_classes10): super().__init__() self.head nn.Linear(embed_dim, num_classes) self.embed_dim embed_dim def forward(self, x): # x 是 ViT 输出的 cls_token形状 [B, 256] return self.head(x)我习惯把 ViT 主干和分类头分开写这样换数据集时只要改num_classes不需要动主干。训练时的 loss 用nn.CrossEntropyLoss()它内部会做 softmax所以模型最后一层不要额外加 softmax。验证时算pred.argmax(dim1)和target的相等率就可以了。这里唯一要注意的是embed_dim与head的输入维度必须一致改embed_dim时漏改分类头是最常见的低级错误通常在第一次 forward 就会报维度错误反倒容易发现。3. 训练与验证跑通 CIFAR10 的完整命令与参数设置3.1 数据加载CIFAR10 的自动下载与标准化这份源码默认使用torchvision.datasets.CIFAR10首次运行会自动下载数据集。由于服务器在国外国内网络下载可能很慢甚至失败建议先手动下载cifar-10-python.tar.gz放到./data目录下torchvision 检测到文件存在后就不会重复下载。数据加载部分一般长这样import torch from torchvision import datasets, transforms 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_test transforms.Compose([ transforms.ToTensor(), transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2470, 0.2435, 0.2616)), ]) train_set datasets.CIFAR10(root./data, trainTrue, downloadTrue, transformtransform_train) test_set datasets.CIFAR10(root./data, trainFalse, downloadTrue, transformtransform_test) train_loader torch.utils.data.DataLoader(train_set, batch_size128, shuffleTrue, num_workers4) test_loader torch.utils.data.DataLoader(test_set, batch_size128, shuffleFalse, num_workers4)注意Normalize的三个均值和标准差是 CIFAR10 的官方统计值不是瞎猜的。如果你换了数据集一定要重新计算均值标准差否则训练起来 loss 异常高且收敛极慢。RandomCrop(32, padding4)是小图上的经典增强padding 的 4 像素默认补零在 32×32 上等价于先放大到 40×40 再随机裁剪回 32×32能有效增加位置多样性。3.2 训练循环优化器、学习率与轮数设置ViT 在 CIFAR10 上一般建议用AdamW初始学习率 1e-3 到 3e-3配合余弦退火。原版 ViT 用的是超大 batch 加 Adam但那是针对 ImageNet 级别的数据量在 CIFAR10 这种 5 万张图的规模上batch size 128 加 60 个 epoch 足够。训练循环的骨架如下import torch.nn as nn from torch.optim import AdamW model ViTClassifier(...) criterion nn.CrossEntropyLoss() optimizer AdamW(model.parameters(), lr3e-3, weight_decay5e-2) scheduler torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max60) for epoch in range(60): model.train() total_loss 0 correct 0 total 0 for images, labels in train_loader: images, labels images.cuda(), labels.cuda() outputs model(images) loss criterion(outputs, labels) optimizer.zero_grad() loss.backward() optimizer.step() total_loss loss.item() * images.size(0) _, pred outputs.max(1) correct pred.eq(labels).sum().item() total labels.size(0) scheduler.step() train_acc 100.0 * correct / total print(fEpoch {epoch1:02d} loss {total_loss/total:.4f} train_acc {train_acc:.2f}%)这里的几个参数你需要根据自己的显卡调batch_size128需要至少 6GB 显存embed_dim256, depth6 的情况下。如果显存不够降到 64 或 32同时学习率最好按比例降低常见做法是 batch 减半、lr 也减半。weight_decay5e-2是 ViT 的常用值比 CNN 通常用的 1e-4 大得多因为 Transformer 更容易过拟合。T_max60必须等于训练总 epoch 数否则余弦退火会在中途终止导致学习率不按预期降到零。验证循环一样简单但注意要写model.eval()和torch.no_grad()否则 BN 和 dropout 状态不对验证准确率会偏低而且会白白占显存。3.3 验证脚本不只看准确率还要看 loss验证不只是算准确率。我习惯在验证时把每个类别的准确率也统计出来特别是用 CIFAR10 这类平衡数据集时总准确率看不出模型在哪些类别上偏弱。model.eval() class_correct [0] * 10 class_total [0] * 10 with torch.no_grad(): for images, labels in test_loader: images, labels images.cuda(), labels.cuda() outputs model(images) _, pred outputs.max(1) correct pred.eq(labels) for i in range(labels.size(0)): label labels[i].item() class_correct[label] correct[i].item() class_total[label] 1 classes [airplane, automobile, bird, cat, deer, dog, frog, horse, ship, truck] for i in range(10): print(f{classes[i]:10s} accuracy: {100.0 * class_correct[i] / class_total[i]:.2f}%)这里通过pred.eq(labels)得到布尔张量再逐样本统计。注意correct[i].item()要把 tensor 转成标量否则累加的是张量最后的除法会报类型错误。分类别统计能让你看到猫和狗这两个类别的准确率大概率低于其它类这不是 bug而是这两个类在 32×32 分辨率下本身就难分属于正常现象。3.4 保存与加载模型别只存权重要存整个训练状态训练完后保存模型有两种常见做法只存权重或者存 checkpoint。对于要做毕设答辩的同学建议保存 checkpoint因为你后续可能还要恢复训练或者微调。torch.save({ model_state_dict: model.state_dict(), optimizer_state_dict: optimizer.state_dict(), epoch: epoch, train_acc: train_acc, }, vit_cifar10_last.pth)加载时注意如果模型类定义在当前文件里直接torch.load然后load_state_dict就行。但如果你改动了模型结构比如把embed_dim从 256 改成 128旧权重会因为形状不匹配而无法加载。这也是很多人拿预训练权重微调时翻车的原因——不匹配的 shape 会报size mismatch for pos_embed解决办法是加载时忽略不匹配的 key或者只加载主干部分。4. 避坑指南ViT 在小数据集上最容易翻车的四个坑4.1 损失不降、准确率在 10% 附近徘徊现象训练了十几个 epoch训练 loss 一直维持在 2.3 左右验证准确率始终在 10% 上下和随机猜测没区别。原因最常见的是学习率太大或者模型输入没有归一化。CIFAR10 的像素值如果不做ToTensor归一化直接以 0-255 的整数喂给模型梯度会爆炸loss 无法收敛。另一个原因是pos_embed拼错了位置但那种情况通常会在第一轮就报维度错误不会静默失败。解决先检查数据管道。把transform里的ToTensor()和Normalize加上确认数据范围是 [-1, 1] 附近。然后把学习率降到 1e-3甚至 1e-4 测试一遍。如果 loss 能在前 5 个 epoch 明显下降说明之前就是学习率问题。我一般会先用一个很小的子集比如 100 张图过拟合到 100% 准确率如果过拟合都做不到问题一定在模型或数据本身而不是训练策略。4.2 训练准确率很高但验证准确率低现象训练集准确率能到 98%验证集只有 70%差距非常大。原因ViT 在 CIFAR10 这种规模的数据集上非常容易过拟合因为参数量远超数据量。Patch 数量少64 个 token意味着模型有大量注意力头可以记忆训练样本尤其在训练轮数超过 100 的情况下。解决把weight_decay提高到 0.1或者加dropout0.2在 Encoder 输出和分类头之间。另外数据增强不能只靠RandomCrop和RandomFlip可以试试Cutout随机遮挡一块正方形区域或者Mixup。我实测下来Mixup的 alpha 取 0.2 时验证准确率能提升 12 个百分点但训练收敛会变慢需要多跑 20 个 epoch 才能看到效果。4.3 换 GPU 之后结果对不上现象同一份代码在 A 机器上训练能到 85% 准确率换到 B 机器后复现只有 80%怎么都拉不齐。原因CIFAR10 数据加载的num_workers不同会导致数据打乱顺序不同但这不是主因。真正的问题是 PyTorch 版本差异尤其是nn.TransformerEncoderLayer在 1.8 和 2.0 之间的GELU实现细节不完全一致导致数值误差积累。另外如果 B 机器上显存更小batch size 被迫调小模型更新步数变多在固定 epoch 数下效果自然不同。解决复现时保证torch.__version__一致设置torch.manual_seed(0)并且把DataLoader的generator固定下来。更稳的做法是在训练脚本里加一个环境检测打印 cuda 版本和 cuDNN 版本对不上就先更新环境。我跟别人联调代码时第一件事就是对比两边的torch.__version__和torch.backends.cudnn.benchmark设置这个开关默认是 False但有些机器上用户会手动打开反而可能导致结果不稳定。4.4 显存不足但不想降低精度现象batch size 128 直接 OOM报 TypeError 或者 CUDA error: out of memory。原因embed_dim256时ViT 在 32×32 上的显存消耗其实不大更大的开销来自 AdamW 的动量状态每个参数要额外存两份副本。如果还开了torch.compile或者 mixed precision显存占用会更难预估。解决先检查是不是num_workers开太多导致内存溢出而不是显存溢出然后把 batch size 降到 64 或 48。如果不想降 batch可以用torch.utils.checkpoint对 Encoder 层的激活做梯度检查点以时间换显存。还有一个隐藏优化把pos_embed从nn.Parameter改成torch.register_buffer因为它不需要梯度参数量虽然只有 65×256但省一点是一点。当然最直接的做法是model torch.nn.DataParallel(model) # 多卡并行注意卡数分布单卡用户就别加这行了DataParallel在小模型上反而会因为进程调度增加额外显存开销。5. 把 CIFAR10 换成自己的数据改动点与验证指标5.1 自定义数据集的目录结构与 Dataset 类如果你的数据集是标准结构每个类别一个文件夹不需要写复杂的 Dataset 类直接用ImageFolder就行。但要注意ImageFolder默认按文件夹名称排序来分配类别索引如果你的文件夹顺序变了同一个类别的索引就变了回传结果会和模型输出对不上。常见的做法是这样from torchvision.datasets import ImageFolder dataset ImageFolder(root./mydata/train, transformtransform_train) print(dataset.class_to_idx) # 确认类别映射class_to_idx是一个字典你得在训练前打印并保存下来尤其换机器后要检查。我见过有人把第一版模型训完了后来重新整理文件夹类别顺序变了导致所有预测结果错位准确率掉到 30% 才发现是标签问题。5.2 修改模型输出维度与训练参数把num_classes从 10 改成你的实际类别数。这里有几个连带改动分类头nn.Linear(embed_dim, num_classes)的输出维度。如果类别数很少比如 2 类建议把num_heads从 8 降到 4depth从 6 降到 4防止过拟合。数据集的图像尺寸如果不是 32×32需要修改patch_embed的image_size同时重新计算pos_embed的尺寸。我一般会在__init__里写一个检查assert self.patch_embed.patch_num self.pos_embed.size(-2) - 1这个断言的价值在于你改image_size或patch_size后如果忘了动pos_embed运行前第一帧就会报错而不是训练到一半维度爆炸。5.3 验证集划分与指标选择CIFAR10 官方有标准测试集但自己的数据集往往只有一个文件夹。这时最好用train_test_split或者Subset划出 10%20% 的验证集而且要做分层划分保证每个类别在验证集里的比例和训练集一致。from sklearn.model_selection import StratifiedShuffleSplit targets [s[1] for s in dataset.samples] split StratifiedShuffleSplit(n_splits1, test_size0.2, random_state42) train_idx, val_idx next(split.split(range(len(dataset)), targets)) train_set torch.utils.data.Subset(dataset, train_idx) val_set torch.utils.data.Subset(dataset, val_idx)这里用StratifiedShuffleSplit而不是普通train_test_split是因为如果你的数据类别不平衡随机切分可能导致某一类在验证集中数量极少甚至缺失准确率统计完全失真。验证时除了准确率建议同时计算每一类的精确率和召回率。对于小数据集我一般只看 top-1 准确率和 loss 曲线如果两者的趋势背离——loss 在降但准确率不升——说明模型在拟合训练集的同时也在输出概率震荡这时候就需要降低学习率或者增加dropout。5.4 数据量太少的补救方案从预训练权重初始化如果你的数据只有几百张图从零训练 ViT 基本不可能收敛。两个可行方案用 ImageNet 预训练的 ViT 权重做初始化然后冻结前三层 Encoder只训练后面几层和分类头。这类预训练权重在timm库里有但注意输入尺寸要是 224×224你需要把图像 resize 到 224或者调整 patch embed 的输入通道。用 CIFAR10 上训练好的权重作为起点做迁移学习。虽然类别不同但低层特征边缘、纹理是通用的。加载时忽略分类头那层然后微调整个网络。我常用第二种因为 CIFAR10 的 ViT 权重比较小加载快而且源域和目标域都是小图特征分布偏差不大。pretrained torch.load(vit_cifar10_last.pth) model.load_state_dict(pretrained[model_state_dict], strictFalse)strictFalse会忽略分类头head.weight和head.bias的 shape 不匹配只加载主干部分。加载后把model.head nn.Linear(embed_dim, new_classes)替换掉新分类头。注意如果你换了embed_dim这一步也没用因为所有权重都不匹配。所以迁移时必须保持embed_dim和depth不变。6. 把验证集当调试工具用准确率波动定位问题的一个实用习惯我训练 ViT 时有个习惯每个 epoch 结束不仅打印验证准确率还会额外打印一个“近 5 个 epoch 的验证准确率标准差”。如果标准差连续保持在 0.8 个百分点以上说明模型在 60 epoch 的学习率阶段内输出不稳定通常意味着学习率过大或者数据增强太强。这个数字是一个很灵敏的探针比单纯看曲线更早暴露问题。具体做法是维护一个长度为 5 的列表存最近 5 个 epoch 的验证准确率然后计算np.stdimport numpy as np recent_accs [] for epoch in range(epochs): val_acc validate(model, val_loader) recent_accs.append(val_acc) if len(recent_accs) 5: recent_accs.pop(0) if len(recent_accs) 5: std np.std(recent_accs) print(frecent acc std: {std:.3f}) if std 0.8: optimizer.param_groups[0][lr] * 0.5 print(lr decay triggered by std)这个技巧来自我踩过的坑之前训练一个五分类的小数据集curriculum 一直显示准确率在 78% 上下震荡我以为是模型容量不够加宽了网络反而更差。后来才发现是学习率一直在高值区间没有降下来曲线像锯齿一样。从那以后我每次跑 ViT 小数据训练都会把这个标准差当成一个硬指标一旦超过 0.8 就手动把学习率减半而不是傻等到 cosine schedule 生效。另外验证集在训练中途的作用不是“判定好坏”而是“定位过拟合”。我习惯在训练到一半时用验证集做一次混淆矩阵可视化如果发现某两个类别的混淆非常严重优先去查数据集里这两个类的样本是不是存在大量标注错误。CIFAR10 上猫和狗的标签噪声就是有名的例子我多次看到训练曲线的尾段猫狗互相吞噬这不是模型问题是标注问题。这份源码给我的最大启发是ViT 做小图分类时真正决定效果的往往不是网络有多深而是 patch_size 与数据增强的配合。把patch_size从 4 改成 8准确率立刻下降 10 个点把RandomCrop去掉准确率又会掉 3 个点左右。所以如果你拿到这份源码先不用急着改架构老老实实把默认参数跑一遍记录下自己的验证准确率曲线然后试着改一个变量——比如patch_size或dropout——再跑一遍对比两条曲线。这种对比实验的意义远大于盲目调参。最后给你一个可以直接用的检查清单跑通后先确认训练 loss 在前 10 个 epoch 内能降到 1.0 以下验证准确率在 50 个 epoch 内至少达到 80%如果达不到优先怀疑学习率和pos_embed重置。这份源码帮你把路铺好了接下来的调试过程才是真正值钱的部分。希望这组笔记能帮你少走几步弯路。本文还有配套的精品资源点击获取
