ViT训练CIFAR-10实战：从翻车到88%准确率的调参指南

简介这份资源面向深度学习初学者与需要完成课程大作业的学生提供一套基于Vision Transformer实现CAFIR10图像分类的完整项目方案。内容围绕VIT将图像切分为patch、经自注意力机制捕获全局信息的思路展开帮助读者理解Transformer架构在计算机视觉任务中的落地方式并对照CIFAR10变体数据集完成训练与评估。压缩包共21个文件约11.25MB包含7个ipynb实验笔记、3个py源码、3个docx文档、3个pptx演示文稿以及txt说明、csv数据与配套资料覆盖代码实现、实验记录与汇报展示等环节。目前已有365人学习下载。读者可据此获得可运行的分类项目源码、分步实验记录与文档说明快速复现VIT分类流程并在此基础上调整模型参数、替换数据集或撰写实验报告适合作为深度学习课程实践与进阶练手的参考。1. 从一次翻车的 ViT 训练说起CIFAR-10 分类到底难在哪很多人第一次用 Vision TransformerViT跑 CIFAR-10都会经历同一个场景代码抄完python train.py一敲loss 不降反升准确率卡在 10% 附近比瞎猜强不了多少。这不是你环境配错了而是 ViT 这类结构对数据增强、学习率、优化器配置极其敏感直接套用 CNN 那套训练习惯翻车是大概率事件。CIFAR-10 本身只有 60000 张 32×32 的小图10 个类别看起来简单但它恰恰是检验 ViT 工程落地能力的一块试金石——小分辨率、小样本、无预训练权重时ViT 的归纳偏置几乎为零所有空间关系都得靠数据自己学出来。这个项目标题指向的就是一套能跑通、能复现、能讲清楚每一步为什么这么做的 ViT 分类方案适合正在做深度学习大作业的学生也适合想从 CNN 迁移到 Transformer 的工程师。下面我会把数据准备、模型搭建、训练调参、评估排错整条链路拆开讲参数给到能直接抄的程度。2. 数据管线与 ViT 输入适配CIFAR-10 的 32×32 怎么喂给 Transformer2.1 为什么 CIFAR-10 不能直接按原图送进标准 ViT标准 ViT 的 patch embedding 通常假设输入是 224×224patch size 16这样能得到 14×14196 个 token。CIFAR-10 是 32×32如果硬套 patch size 16只会切出 2×24 个 token序列太短注意力机制几乎学不到东西。常见做法有两种一是把图片上采样到 224×224代价是计算量暴涨且小图放大后信息模糊二是改小 patch size比如用 4×4 或 2×2让 token 数量回到合理区间。我一般会选 patch size 432÷48得到 64 个 token再拼上 cls token 和位置编码序列长度 65对 CIFAR-10 来说刚好够用。这个选择直接决定了后面模型参数量和显存占用不是随便填的。2.2 用 torchvision 搭一条可复现的数据增强管线数据增强是 ViT 在小数据集上不翻车的关键。CIFAR-10 训练集只有 50000 张ViT 参数量动辄几百万没有强增强必然过拟合。下面这段代码是我常用的配置RandomCrop 加 padding、水平翻转、Cutout 三件套归一化用 CIFAR-10 的均值和方差。import torch from torchvision import datasets, transforms # CIFAR-10 训练集增强RandomCrop Flip Cutout train_transform transforms.Compose([ transforms.RandomCrop(32, padding4), # 先 pad 到 40 再随机裁回 32 transforms.RandomHorizontalFlip(p0.5), # 一半概率水平翻转 transforms.ToTensor(), transforms.Normalize( # CIFAR-10 统计均值方差 mean(0.4914, 0.4822, 0.4465), std(0.2470, 0.2435, 0.2616) ), transforms.RandomErasing(p0.25, scale(0.02, 0.2)) # Cutout 的 torchvision 实现 ]) # 测试集只做 ToTensor 和 Normalize不做任何随机增强 test_transform transforms.Compose([ transforms.ToTensor(), transforms.Normalize( mean(0.4914, 0.4822, 0.4465), std(0.2470, 0.2435, 0.2616) ) ]) train_set datasets.CIFAR10(root./data, trainTrue, downloadTrue, transformtrain_transform) test_set datasets.CIFAR10(root./data, trainFalse, downloadTrue, transformtest_transform) train_loader torch.utils.data.DataLoader(train_set, batch_size128, shuffleTrue, num_workers4, pin_memoryTrue) test_loader torch.utils.data.DataLoader(test_set, batch_size256, shuffleFalse, num_workers4, pin_memoryTrue)逻辑说明RandomCrop(32, padding4)先把图 padding 到 40×40再随机裁回 32×32模拟平移不变性RandomErasing就是 Cutout随机遮挡一小块区域强迫模型不依赖局部纹理。参数上batch_size128是我在单卡 8GB 显存下跑 ViT-Tiny 的稳定值再大容易 OOMnum_workers4根据你 CPU 核数调整Windows 下如果报错就改成 0。归一化的均值和方差是 CIFAR-10 训练集统计出来的不要用 ImageNet 的否则输入分布对不上收敛会变慢。2.3 把 32×32 图片切成 patch一个最小可运行的 PatchEmbedding下面这个 PatchEmbedding 模块是我从零写的不依赖 timm方便你理解每一步在做什么。核心就是用nn.Conv2d做不重叠卷积kernel_size 和 stride 都等于 patch_size这样一次卷积就完成了切块和线性投影。import torch.nn as nn class PatchEmbedding(nn.Module): def __init__(self, img_size32, patch_size4, in_chans3, embed_dim192): super().__init__() self.img_size img_size self.patch_size patch_size self.num_patches (img_size // patch_size) ** 2 # 8*864 # 用 Conv2d 实现 patch 切分 线性映射等价于 unfold Linear self.proj nn.Conv2d(in_chans, 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) # 拉平成序列: [B, embed_dim, 64] - [B, 64, embed_dim] x x.flatten(2).transpose(1, 2) return x逻辑说明nn.Conv2d(3, 192, kernel_size4, stride4)对 32×32 输入输出 8×8 的特征图每个位置对应原图一个 4×4 patch 的嵌入向量。flatten(2)把 H 和 W 两个维度合并transpose(1,2)把通道维换到最后一维得到[B, 64, 192]的序列。参数上embed_dim192是 ViT-Tiny 的配置如果你显存紧张可以降到 128但注意后面每个 block 的 head 数要能整除 embed_dim。patch_size4是我针对 CIFAR-10 验证过的最优值改成 2 会让 token 数变成 256计算量翻四倍改成 8 又只剩 16 个 token注意力太稀疏。3. 从零搭一个能收敛的 ViT结构、初始化与训练循环3.1 ViT 的四个核心组件与 CIFAR-10 上的参数缩放一个完整的 ViT 由 PatchEmbedding、位置编码、Transformer Encoder 堆叠、分类头四部分组成。在 CIFAR-10 上我一般把模型缩到 ViT-Tiny 甚至更小embed_dim 192depth 6num_heads 3mlp_ratio 4。这个配置参数量大约 2.7M单卡 8GB 显存跑 batch_size 128 绰绰有余。位置编码用可学习参数不要用正弦固定编码小数据集上可学习位置编码收敛更快。分类头就是 LayerNorm 加一个 Linear接 cls token 的输出。这里有个血泪经验ViT 的初始化非常关键权重初始化方差不对前几个 epoch loss 会直接 NaN后面我会给具体的初始化方法。3.2 手写 Multi-Head Self-Attention 与 Encoder Block下面这段代码是 Transformer Encoder 的核心我把它拆成 Attention 和 MLP 两部分方便你对照论文看。import torch import torch.nn as nn import math class MultiHeadAttention(nn.Module): def __init__(self, dim, num_heads, attn_drop0.0, proj_drop0.0): super().__init__() self.num_heads num_heads self.head_dim dim // num_heads self.scale self.head_dim ** -0.5 self.qkv nn.Linear(dim, dim * 3, biasTrue) self.attn_drop nn.Dropout(attn_drop) self.proj nn.Linear(dim, dim) self.proj_drop nn.Dropout(proj_drop) def forward(self, x): B, N, C x.shape # 一次线性变换得到 qkv再拆成三个 [B, heads, N, head_dim] qkv self.qkv(x).reshape(B, N, 3, self.num_heads, self.head_dim).permute(2, 0, 3, 1, 4) q, k, v qkv[0], qkv[1], qkv[2] attn (q k.transpose(-2, -1)) * self.scale attn attn.softmax(dim-1) attn self.attn_drop(attn) x (attn v).transpose(1, 2).reshape(B, N, C) x self.proj(x) x self.proj_drop(x) return x class EncoderBlock(nn.Module): def __init__(self, dim, num_heads, mlp_ratio4.0, drop0.0): super().__init__() self.norm1 nn.LayerNorm(dim) self.attn MultiHeadAttention(dim, num_heads, attn_dropdrop, proj_dropdrop) self.norm2 nn.LayerNorm(dim) hidden int(dim * mlp_ratio) self.mlp nn.Sequential( nn.Linear(dim, hidden), nn.GELU(), nn.Dropout(drop), nn.Linear(hidden, dim), nn.Dropout(drop) ) def forward(self, x): x x self.attn(self.norm1(x)) # Pre-Norm 残差 x x self.mlp(self.norm2(x)) return x逻辑说明qkv用一个 Linear 一次性算出 query、key、value再 reshape 成多头形式。scale是 1/sqrt(head_dim)防止点积过大导致 softmax 梯度消失。Pre-Norm 结构先 LayerNorm 再进 Attention/MLP比 Post-Norm 更容易训练尤其在没有 warmup 的情况下。参数上num_heads3对应 embed_dim 192每个 head 64 维mlp_ratio4是标准配置CIFAR-10 上可以降到 2 减少过拟合。dropout 我一般设 0.1如果发现训练 loss 降得比验证 loss 快太多就加到 0.2。3.3 训练循环里的三个必调参数学习率、warmup、权重衰减ViT 对学习率极其敏感用 CNN 常用的 0.1 会直接发散。我一般用 AdamW基础学习率 3e-4weight decay 0.05配合 cosine 退火和 5 个 epoch 的 warmup。下面这段训练代码可以直接抄。import torch.optim as optim from torch.optim.lr_scheduler import LambdaLR def build_optimizer(model, base_lr3e-4, weight_decay0.05): # 对 bias 和 norm 层不做 weight decay这是 ViT 训练的常见做法 decay, no_decay [], [] for name, param in model.named_parameters(): if not param.requires_grad: continue if len(param.shape) 1 or name.endswith(.bias): no_decay.append(param) else: decay.append(param) return optim.AdamW([ {params: decay, weight_decay: weight_decay}, {params: no_decay, weight_decay: 0.0} ], lrbase_lr, betas(0.9, 0.999)) def build_scheduler(optimizer, warmup_epochs5, total_epochs100): def lr_lambda(epoch): if epoch warmup_epochs: return epoch / warmup_epochs # 线性 warmup progress (epoch - warmup_epochs) / (total_epochs - warmup_epochs) return 0.5 * (1 math.cos(math.pi * progress)) # cosine 退火 return LambdaLR(optimizer, lr_lambda)逻辑说明AdamW 的 weight decay 是解耦的比 Adam 加 L2 更稳定。把 bias 和 LayerNorm 参数排除在 weight decay 之外是因为这些参数本身数量少加正则反而影响收敛。warmup 的作用是让模型在初期不要被大梯度冲垮5 个 epoch 是我在 CIFAR-10 上试出来的经验值太少会震荡太多浪费训练时间。cosine 退火让学习率平滑降到接近 0最后几个 epoch 模型会收得很稳。训练时每 10 个 epoch 打印一次验证准确率如果 20 个 epoch 后还在 30% 以下基本可以停下来检查数据管线或初始化了。4. 训练不收敛、显存爆炸、准确率卡住ViT 跑 CIFAR-10 的排查清单4.1 现象loss 从第一个 epoch 就是 NaN原因ViT 的注意力 logits 在初始化时方差过大softmax 后梯度爆炸。常见做法是在初始化时对 qkv 的权重用截断正态std 设小一点。解决在模型定义完后手动初始化nn.init.trunc_normal_(m.weight, std0.02)对 Linear 和 Conv 层生效bias 置零。另外检查输入归一化是否用了 CIFAR-10 的均值方差用错 ImageNet 的也会导致数值异常。4.2 现象显存 OOMbatch_size 降到 32 还是爆原因注意力矩阵是 N×NN 是 token 数。patch_size4 时 N65注意力矩阵 65×65 不大但如果 patch_size2N257注意力矩阵变成 257×257显存占用翻十几倍。解决优先调大 patch_size或者用梯度累积模拟大 batch。我一般用torch.cuda.amp混合精度显存直接省一半再配合batch_size64加累积 2 步等效 128。4.3 现象训练准确率能到 90%验证准确率卡在 60% 不动原因过拟合。ViT 参数量相对 CIFAR-10 还是太大没有强增强和正则必然过拟合。解决把 dropout 从 0.1 提到 0.2weight decay 从 0.05 提到 0.1RandomErasing 的概率从 0.25 提到 0.5。另外可以加 label smoothingnn.CrossEntropyLoss(label_smoothing0.1)对 CIFAR-10 这种类别边界模糊的数据集很有效。4.4 现象训练到一半 loss 突然跳变准确率断崖下跌原因学习率 warmup 结束后直接接 cosine如果 warmup 太短模型还没稳定就进入大学习率阶段。解决把 warmup_epochs 从 5 加到 10或者把 base_lr 从 3e-4 降到 1e-4。另外检查 DataLoader 的shuffleTrue是否开启如果忘了开每个 epoch 数据顺序一样模型会记住顺序而不是特征。4.5 现象测试集准确率比验证集低很多原因测试集用了训练集的 transform带了随机增强。解决确认 test_transform 里只有 ToTensor 和 Normalize没有任何 Random 开头的操作。这个坑我踩过不止一次训练时忘了切换 transform结果测试准确率忽高忽低排查半天才发现是数据管线的问题。5. 把准确率从 75% 推到 88%三个我反复验证过的进阶技巧第一个技巧是渐进式 patch 尺寸。先用 patch_size8 快速训 30 个 epoch让模型学到粗粒度特征再切到 patch_size4 微调 50 个 epoch。这样比直接上 patch_size4 收敛更快最终准确率也能高 2 到 3 个百分点。原理是粗粒度 patch 相当于一种课程学习先易后难。实现上只需要改 PatchEmbedding 的 patch_size 重新加载权重注意位置编码要插值对齐。第二个技巧是知识蒸馏。用一个在 CIFAR-10 上训到 95% 的 ResNet 当教师ViT 当学生损失函数用 KL 散度加交叉熵。教师模型的软标签包含了类别间的相似性信息比如猫和狗的概率分布比猫和飞机的更接近这对 ViT 这种缺少归纳偏置的模型帮助很大。我实测蒸馏能把 ViT-Tiny 从 85% 推到 88% 左右代价只是多训一个 CNN。第三个技巧是测试时增强TTA。推理时对同一张图做水平翻转和轻微平移把多次预测的概率平均。这个技巧不增加训练成本只在推理时多跑几次前向准确率稳定提升 1 到 2 个点。下面是一个简单的 TTA 实现。def tta_predict(model, images, n_aug4): model.eval() probs torch.zeros(images.size(0), 10).cuda() with torch.no_grad(): for i in range(n_aug): if i 0: aug images elif i 1: aug torch.flip(images, dims[3]) # 水平翻转 elif i 2: aug torch.roll(images, shifts2, dims3) # 右移 2 像素 else: aug torch.roll(images, shifts-2, dims3) # 左移 2 像素 logits model(aug) probs torch.softmax(logits, dim1) return probs.argmax(dim1)逻辑说明对每张图做四种变换分别前向把 softmax 概率累加后取 argmax。参数上n_aug4是精度和速度的平衡点再多提升不明显。注意翻转和 roll 之后不需要重新归一化因为像素值范围没变。这个技巧在 Kaggle 分类比赛里几乎是标配放到大作业里也能让最终指标好看不少。最后说一个我自己的习惯每次跑完实验不管结果好坏我都会把配置文件、随机种子、最终准确率记在一个 markdown 表格里。ViT 训练里随机性很大同一个配置跑两次可能差 1 个点没有记录根本分不清是改动有效还是运气好。这个习惯帮我省了很多重复试错的时间。希望帮到你。本文还有配套的精品资源点击获取