Python+ViT实战：CIFAR-10十分类从数据加载到90%精度

简介这份资源面向深度学习初学者与课程实践者提供一套基于Vision Transformer完成CIFAR10图像分类的完整大作业方案帮助读者理解如何用Python将Transformer架构迁移到计算机视觉任务中掌握图像分块、自注意力机制与分类头设计的核心思路。压缩包共21个文件约11.25MB包含7个ipynb实验笔记、3个py源码、3份docx文档、3份pptx汇报材料以及csv数据与txt说明覆盖从数据加载、模型搭建到训练评估的全流程。资源还附带手写数字识别、机器翻译、LSTM自动写诗等同类作业参考便于横向对比不同模型的应用方式。目前已有365人学习适合需要完成课程设计、快速复现VIT分类实验并整理实验报告的学生参考文档与代码配合使用可减少环境配置与调试成本。1. 从一次翻车的 ViT 训练说起CIFAR-10 分类到底难在哪很多人第一次拿 ViT 做 CIFAR-10信心满满地把patch_size设成 16结果训练 20 个 epoch 后准确率卡在 60% 上下连一个调好的 ResNet-18 都打不过。这不是代码写错了而是 ViT 原论文的配置是给 ImageNet 这种 224×224 大图设计的直接搬到 32×32 的 CIFAR-10 上patch 数量、感受野、位置编码全都对不上。这个标题讲的就是用 Python 和 Vision Transformer 完成 CIFAR-10 十分类从数据加载、模型改造、训练调参到评估导出走通一条能复现的路径。适合正在做深度学习大作业的学生、想从 CNN 切换到 Transformer 的工程师以及需要一份可运行源码加文档参考的从业者。CIFAR-10 本身只有 6 万张 32×32 彩色图10 个类别看似简单但它是检验 ViT 小图适配能力的一块试金石也是理解 patch embedding、注意力机制和分类头设计的最小闭环。2. ViT 做 CIFAR-10 的模型改造patch 怎么切、位置编码怎么加2.1 为什么原版 ViT 在 32×32 图上会水土不服原版 ViT 的输入是 224×224patch_size 取 16得到 14×14196 个 patch加上 cls token 一共 197 个 token。CIFAR-10 的图只有 32×32如果还用 patch_size16只能切出 2×24 个 patch序列长度太短注意力机制几乎退化成全连接根本学不到空间关系。常见做法是把 patch_size 降到 4这样 32÷48得到 8×864 个 patch序列长度变成 65和原版 197 相比虽然短但已经能覆盖足够的空间位置。另一个坑是位置编码原版用 1D 可学习位置编码长度固定为 197直接加载预训练权重会维度不匹配。我一般会重新初始化位置编码或者用插值的方式把 197 个位置编码降采样到 65 个但插值实现起来麻烦重新初始化更省事反正 CIFAR-10 从头训练也能收敛。2.2 用 PyTorch 搭建一个适配 CIFAR-10 的 ViT下面这段代码定义了一个精简版 ViTpatch_size4embed_dim256depth6num_heads8mlp_ratio4dropout0.1。这个配置在单张 8GB 显存的卡上就能跑batch_size 可以设到 128。import torch import torch.nn as nn from einops import rearrange class PatchEmbedding(nn.Module): def __init__(self, img_size32, patch_size4, in_chans3, embed_dim256): super().__init__() self.patch_size patch_size self.num_patches (img_size // patch_size) ** 2 # 64 # 用卷积实现切 patch 和线性投影等价于 unfold linear self.proj nn.Conv2d(in_chans, embed_dim, kernel_sizepatch_size, stridepatch_size) def forward(self, x): x self.proj(x) # (B, embed_dim, 8, 8) x rearrange(x, b c h w - b (h w) c) # (B, 64, embed_dim) return x class ViTForCIFAR10(nn.Module): def __init__(self, img_size32, patch_size4, in_chans3, num_classes10, embed_dim256, depth6, num_heads8, mlp_ratio4, dropout0.1): super().__init__() self.patch_embed PatchEmbedding(img_size, patch_size, in_chans, embed_dim) num_patches self.patch_embed.num_patches self.cls_token nn.Parameter(torch.zeros(1, 1, embed_dim)) self.pos_embed nn.Parameter(torch.zeros(1, num_patches 1, embed_dim)) self.pos_drop nn.Dropout(dropout) # Transformer Encoder 层 encoder_layer nn.TransformerEncoderLayer( d_modelembed_dim, nheadnum_heads, dim_feedforwardembed_dim * mlp_ratio, dropoutdropout, activationgelu, batch_firstTrue, norm_firstTrue ) self.encoder nn.TransformerEncoder(encoder_layer, num_layersdepth) self.norm nn.LayerNorm(embed_dim) self.head nn.Linear(embed_dim, num_classes) def forward(self, x): B x.shape[0] x self.patch_embed(x) # (B, 64, embed_dim) cls_tokens self.cls_token.expand(B, -1, -1) x torch.cat([cls_tokens, x], dim1) # (B, 65, embed_dim) x x self.pos_embed x self.pos_drop(x) x self.encoder(x) x self.norm(x[:, 0]) # 取 cls token 的输出 return self.head(x)逻辑说明PatchEmbedding用nn.Conv2d一步完成切块和线性映射比手动 unfold 再 matmul 更简洁也更快。cls_token和pos_embed都是可学习参数pos_embed的长度是num_patches 1正好对应 64 个 patch 加 1 个 cls token。nn.TransformerEncoderLayer里norm_firstTrue表示 Pre-LN训练更稳定batch_firstTrue让输入维度是(B, seq_len, dim)省去转置操作。参数说明embed_dim256是每个 patch 映射后的向量维度depth6是 Transformer 层数num_heads8要求embed_dim能被num_heads整除mlp_ratio4表示 FFN 中间层维度是embed_dim * 4。如果显存不够优先降embed_dim到 192 或 128再降depth到 4。2.3 数据增强和归一化参数怎么设CIFAR-10 的均值和标准差是(0.4914, 0.4822, 0.4465)和(0.2470, 0.2435, 0.2616)不要用 ImageNet 的(0.485, 0.456, 0.406)否则输入分布偏移收敛会变慢。训练时用 RandomCrop(32, padding4) 和 RandomHorizontalFlip测试时只做 ToTensor 和 Normalize。下面这段是 DataLoader 的配置from torchvision import datasets, transforms from torch.utils.data import DataLoader train_transform transforms.Compose([ transforms.RandomCrop(32, padding4), transforms.RandomHorizontalFlip(), transforms.ToTensor(), transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2470, 0.2435, 0.2616)) ]) test_transform 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, transformtrain_transform) test_set datasets.CIFAR10(root./data, trainFalse, downloadTrue, transformtest_transform) train_loader DataLoader(train_set, batch_size128, shuffleTrue, num_workers4, pin_memoryTrue) test_loader DataLoader(test_set, batch_size256, shuffleFalse, num_workers4, pin_memoryTrue)RandomCrop(32, padding4)先 padding 到 40×40 再随机裁回 32×32相当于平移增强。num_workers4在 Linux 上没问题Windows 下如果报错就改成 0。pin_memoryTrue配合 GPU 训练能加快数据搬运。3. 训练循环与调参学习率、优化器、混合精度怎么配3.1 优化器选 AdamW 还是 SGDViT 原论文用 AdamWbetas(0.9, 0.999)weight_decay0.05。CIFAR-10 上我试过 SGD momentum0.9收敛更慢最终精度差不多但需要更长的 epoch。AdamW 对学习率更敏感但配合 warmup 和 cosine 衰减通常 100 个 epoch 能到 85% 以上。下面是一个带 warmup 的 cosine 调度器import math from torch.optim import AdamW from torch.optim.lr_scheduler import LambdaLR def get_cosine_schedule_with_warmup(optimizer, warmup_epochs, total_epochs, base_lr, min_lr1e-5): def lr_lambda(epoch): if epoch warmup_epochs: return epoch / warmup_epochs progress (epoch - warmup_epochs) / (total_epochs - warmup_epochs) return min_lr / base_lr (1 - min_lr / base_lr) * 0.5 * (1 math.cos(math.pi * progress)) return LambdaLR(optimizer, lr_lambda) model ViTForCIFAR10().cuda() optimizer AdamW(model.parameters(), lr3e-4, weight_decay0.05) scheduler get_cosine_schedule_with_warmup(optimizer, warmup_epochs10, total_epochs100, base_lr3e-4)base_lr3e-4是 AdamW 在 ViT 上的常用起点如果 loss 震荡就降到 1e-4。warmup_epochs10让学习率从 0 线性升到 3e-4避免初期梯度爆炸。min_lr1e-5是余弦衰减的终点不要设成 0否则后期完全学不动。3.2 混合精度训练和梯度裁剪单卡 8GB 显存跑embed_dim256, depth6, batch_size128基本够用但如果你想加大 batch_size 或 embed_dim可以用 AMP 混合精度。下面这段是训练一个 epoch 的核心逻辑from torch.cuda.amp import autocast, GradScaler scaler GradScaler() criterion nn.CrossEntropyLoss() def train_one_epoch(model, loader, optimizer, scheduler, scaler, device): model.train() total_loss, correct, total 0, 0, 0 for imgs, labels in loader: imgs, labels imgs.to(device), labels.to(device) optimizer.zero_grad() with autocast(): outputs model(imgs) loss criterion(outputs, labels) scaler.scale(loss).backward() scaler.unscale_(optimizer) torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm1.0) scaler.step(optimizer) scaler.update() scheduler.step() total_loss loss.item() * imgs.size(0) _, preds outputs.max(1) correct preds.eq(labels).sum().item() total imgs.size(0) return total_loss / total, correct / totalautocast()自动把部分算子转成 float16节省显存并加速。GradScaler负责放大 loss 避免梯度下溢。clip_grad_norm_把梯度范数限制在 1.0防止个别 batch 梯度爆炸。注意scheduler.step()放在scaler.update()之后每个 batch 都调一次如果按 epoch 调就把这行移到 epoch 循环里。3.3 评估与保存最佳模型评估时关掉梯度用model.eval()和torch.no_grad()。保存模型时同时存state_dict和优化器状态方便断点续训。下面是一个完整的评估函数torch.no_grad() def evaluate(model, loader, device): model.eval() correct, total 0, 0 for imgs, labels in loader: imgs, labels imgs.to(device), labels.to(device) outputs model(imgs) _, preds outputs.max(1) correct preds.eq(labels).sum().item() total imgs.size(0) return correct / total best_acc 0 for epoch in range(100): train_loss, train_acc train_one_epoch(model, train_loader, optimizer, scheduler, scaler, device) val_acc evaluate(model, test_loader, device) if val_acc best_acc: best_acc val_acc torch.save({model: model.state_dict(), optimizer: optimizer.state_dict(), epoch: epoch}, best_vit_cifar10.pth) print(fEpoch {epoch}: train_loss{train_loss:.4f}, train_acc{train_acc:.4f}, val_acc{val_acc:.4f})best_acc记录最高验证精度每次刷新就覆盖保存。torch.save里存一个字典包含模型、优化器和当前 epoch恢复时用torch.load再分别load_state_dict。如果验证精度连续 20 个 epoch 不涨可以提前停但 CIFAR-10 上 ViT 通常到 80 个 epoch 后才开始饱和。4. 避坑与排查CIFAR-10 上跑 ViT 最容易翻车的 5 个地方4.1 现象训练 loss 不降准确率停在 10%原因位置编码维度对不上或者 cls token 没加进去。常见的是pos_embed长度设成了num_patches而不是num_patches 1导致广播时形状错误但 PyTorch 有时不报错只是静默算错。解决打印x.shape和pos_embed.shape确认第二维一致。另外检查cls_token是否在torch.cat时用了expand而不是repeatexpand不复制内存更省显存。4.2 现象验证精度比训练精度高很多原因训练时用了 Dropout 和 RandomCrop验证时没有如果模型欠拟合验证集反而表现更好。但更常见的是数据泄露比如把测试集混进了训练集。解决检查datasets.CIFAR10的train参数训练集用trainTrue测试集用trainFalse。另外确认DataLoader的shuffle在测试集上是False。4.3 现象显存溢出batch_size 降到 32 还是 OOM原因nn.TransformerEncoderLayer默认batch_firstFalse如果没设batch_firstTrue输入会被当成(seq_len, batch, dim)注意力矩阵维度爆炸。解决显式设置batch_firstTrue并检查输入张量形状是(B, seq_len, dim)。另外autocast要配合GradScaler只加autocast不加scaler会导致梯度下溢loss 变成 NaN。4.4 现象训练到一半 loss 突然变成 NaN原因学习率太大或者梯度爆炸。AdamW 的weight_decay0.05在 CIFAR-10 上偏大可以降到 0.01。解决加梯度裁剪clip_grad_norm_(model.parameters(), max_norm1.0)并把base_lr从 3e-4 降到 1e-4。如果还是 NaN检查输入归一化是否用了 CIFAR-10 的均值和标准差用错 ImageNet 的参数会导致输入值域偏移。4.5 现象最终精度只有 75% 左右上不去原因patch_size 太大比如用了 8得到 4×416 个 patch序列太短。或者embed_dim太小比如 128表达能力不够。解决patch_size 降到 4embed_dim升到 256 或 384depth升到 6 或 8。如果显存不够用梯度累积模拟大 batch比如accum_steps4每 4 个 batch 才optimizer.step()一次。5. 进阶技巧用预训练权重和 CutMix 把 CIFAR-10 精度推到 90% 以上如果你不想从头训练可以加载 ImageNet 上预训练的 ViT 权重但位置编码维度不匹配需要插值。PyTorch 的torch.nn.functional.interpolate可以做 1D 插值把 197 个位置编码降采样到 65 个。下面是一个插值函数import torch.nn.functional as F def interpolate_pos_embed(pos_embed, num_patches): # pos_embed: (1, 197, dim) - (1, num_patches1, dim) cls_pos pos_embed[:, :1, :] patch_pos pos_embed[:, 1:, :] # (1, 196, dim) patch_pos patch_pos.reshape(1, 14, 14, -1).permute(0, 3, 1, 2) # (1, dim, 14, 14) patch_pos F.interpolate(patch_pos, size(8, 8), modebilinear, align_cornersFalse) patch_pos patch_pos.permute(0, 2, 3, 1).reshape(1, 64, -1) return torch.cat([cls_pos, patch_pos], dim1)这个函数把 14×14 的 patch 位置编码双线性插值到 8×8再拼上 cls token 的位置编码。加载预训练权重时除了pos_embed和head其他参数直接load_state_dict。head因为类别数从 1000 变成 10需要重新初始化。另一个技巧是 CutMix它把两张图裁剪拼接标签按面积加权。CIFAR-10 上 CutMix 比 MixUp 更稳因为 CutMix 保留了局部空间结构ViT 的注意力机制能更好利用。下面是一个 CutMix 的实现import numpy as np def cutmix_data(x, y, alpha1.0): lam np.random.beta(alpha, alpha) batch_size x.size(0) index torch.randperm(batch_size).to(x.device) bbx1, bby1, bbx2, bby2 rand_bbox(x.size(), lam) x[:, :, bbx1:bbx2, bby1:bby2] x[index, :, bbx1:bbx2, bby1:bby2] lam 1 - ((bbx2 - bbx1) * (bby2 - bby1) / (x.size(-1) * x.size(-2))) y_a, y_b y, y[index] return x, y_a, y_b, lam def rand_bbox(size, lam): W, H size[2], size[3] cut_rat np.sqrt(1. - lam) cut_w, cut_h int(W * cut_rat), int(H * cut_rat) cx, cy np.random.randint(W), np.random.randint(H) bbx1 np.clip(cx - cut_w // 2, 0, W) bby1 np.clip(cy - cut_h // 2, 0, H) bbx2 np.clip(cx cut_w // 2, 0, W) bby2 np.clip(cy cut_h // 2, 0, H) return bbx1, bby1, bbx2, bby2训练时 loss 要改成lam * criterion(outputs, y_a) (1 - lam) * criterion(outputs, y_b)。CutMix 的alpha1.0控制裁剪面积分布CIFAR-10 上 1.0 效果不错太大比如 2.0 会让图变得太碎ViT 反而学不好。我自己的习惯是先用patch_size4, embed_dim256, depth6从头训 100 个 epoch拿到 85% 左右的基线然后加载 ImageNet 预训练权重插值位置编码再训 50 个 epoch精度能到 90% 以上最后加 CutMix 和 label smoothing 0.1再训 30 个 epoch通常能到 92% 左右。整个过程在单张 8GB 卡上大约 6 到 8 小时如果卡更小就把embed_dim降到 192batch_size降到 64时间会拉长但精度掉得不多。希望帮到你。本文还有配套的精品资源点击获取