简介:这份资源面向深度学习入门者与计算机视觉方向的初学者,围绕PyTorch框架与CIFAR-10数据集,提供一套可直接运行的图像识别实践材料,帮助读者理解卷积神经网络从数据加载到模型训练、再到权重复用的完整链路。压缩包共5个文件,以2个Python脚本为核心,分别负责CIFAR-10图片数据的读取与预处理、CNN网络结构定义及训练流程;另含1个已训练好的模型权重文件,可直接加载用于预测或微调,以及1个说明文档和1个数据集元信息文件,整体约7.15MB,结构紧凑、便于快速上手。目前已有53人学习。通过动手运行这些代码,读者能掌握图像标准化与数据增强、损失函数与优化器选择、模型保存与加载等关键环节,并借助现成权重省去从零训练的时间成本,适合作为图像识别与机器学习入门练手、课程实验或项目原型的参考起点。
1. 拆开这个 PyTorch CIFAR-10 图像识别包:它到底能跑出什么结果
如果你手头正好有一个基于PyTorch的CIFAR-10图像识别.zip,别急着双击解压然后对着目录发呆。CIFAR-10 这个数据集在图像识别圈子里算是「Hello World」级别的存在——10 个类别、60000 张 32×32 彩色图、50000 训练 + 10000 测试,标准得不能再标准。但恰恰因为它太标准,网上流传的代码质量参差不齐,有的跑出来准确率 60% 就敢说「完成」,有的连数据增强都没做就硬训。
这个包的核心价值在于:它把「PyTorch 环境搭建 → 数据加载与增强 → 模型定义 → 训练循环 → 评估与推理」这条链路完整地串起来了。适合两类人:一是刚装完 PyTorch、想找个能跑通的项目练手的新手;二是需要快速验证某个 backbone 或训练策略在 CIFAR-10 上表现的老手。我拆过不少类似的包,最怕的就是「代码能跑但结果不可复现」——随机种子没固定、归一化参数写错、验证集划分有泄漏,这些坑后面会一个个说。
2. 环境与数据管线:从 conda 到 DataLoader 的完整链路
2.1 环境搭建:CUDA、cuDNN 与 PyTorch 版本对齐
拿到包第一件事不是pip install -r requirements.txt,而是先确认你的显卡驱动和 CUDA 版本。PyTorch 官方现在推荐用 conda 装,因为 conda 会把 CUDA runtime 和 cuDNN 一起打包,省得你手动配环境变量。我一般这么干:
# 创建独立环境,Python 版本别太新,3.9~3.11 最稳 conda create -n cifar10 python=3.10 -y conda activate cifar10 # 去 PyTorch 官网查对应 CUDA 版本的安装命令,这里以 CUDA 11.8 为例 conda install pytorch torchvision torchaudio pytorch-cuda=11.8 -c pytorch -c nvidia -y # 验证 GPU 是否可用 python -c "import torch; print(torch.__version__, torch.cuda.is_available(), torch.cuda.get_device_name(0))"逻辑说明:pytorch-cuda=11.8这个参数是关键,它决定了 PyTorch 编译时链接的 CUDA 版本。如果你装完发现torch.cuda.is_available()返回False,九成是版本没对齐——要么驱动太老,要么 conda 源里没有对应包。参数上,python=3.10是保守选择,3.12 在某些旧版 torchvision 上会有兼容问题。
提示:如果你用的是 WSL2,记得在 Windows 侧装好显卡驱动,WSL 里不需要再装驱动,直接装 CUDA toolkit 即可。
2.2 数据加载:CIFAR-10 的下载、归一化与增强策略
CIFAR-10 用torchvision.datasets.CIFAR10一行就能下载,但归一化参数别乱填。它的 RGB 三通道均值是(0.4914, 0.4822, 0.4465),标准差是(0.2023, 0.1994, 0.2010)——这是官方统计出来的,用(0.5, 0.5, 0.5)也能跑,但收敛会慢一点。
import torch from torchvision import datasets, transforms from torch.utils.data import DataLoader # 训练集增强:随机裁剪 + 随机水平翻转,这是 CIFAR-10 的标配 train_transform = transforms.Compose([ transforms.RandomCrop(32, padding=4), # 先 pad 4 像素再随机裁回 32×32 transforms.RandomHorizontalFlip(p=0.5), # 一半概率水平翻转 transforms.ToTensor(), # 转成 [0,1] 的 tensor transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)) ]) # 测试集只做 ToTensor + Normalize,不做增强 test_transform = transforms.Compose([ transforms.ToTensor(), transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)) ]) train_set = datasets.CIFAR10(root='./data', train=True, download=True, transform=train_transform) test_set = datasets.CIFAR10(root='./data', train=False, download=True, transform=test_transform) train_loader = DataLoader(train_set, batch_size=128, shuffle=True, num_workers=4, pin_memory=True) test_loader = DataLoader(test_set, batch_size=256, shuffle=False, num_workers=4, pin_memory=True)逻辑说明:RandomCrop(32, padding=4)是 CIFAR-10 的经典增强,先四周补零再随机裁,等价于让物体位置有微小偏移。pin_memory=True在 GPU 训练时能加速 CPU 到 GPU 的数据搬运。num_workers设成 4 是折中值,设太大在 Windows 上容易出多进程问题,设 0 则数据加载会成为瓶颈。
参数上,batch_size=128对 8GB 显存的卡比较友好,如果你用 3090/4090 可以拉到 256 甚至 512,但学习率要相应调大。shuffle=True只对训练集开,测试集必须关,否则评估结果没有意义。
2.3 模型选型:ResNet-18 还是自己搭 CNN
包里的模型定义通常是两种路子:一种是手写一个 3 层卷积 + 全连接的小网络,另一种是直接调torchvision.models.resnet18(pretrained=False)然后改第一层和最后一层。我建议用 ResNet-18,因为它在 CIFAR-10 上不加任何 trick 就能到 93% 以上,而手写小网络往往卡在 75% 左右。
import torch.nn as nn from torchvision import models def build_resnet18(num_classes=10): model = models.resnet18(weights=None) # 不加载 ImageNet 预训练 # CIFAR-10 是 32×32,ImageNet 是 224×224,第一层卷积要改 model.conv1 = nn.Conv2d(3, 64, kernel_size=3, stride=1, padding=1, bias=False) model.maxpool = nn.Identity() # 去掉 maxpool,保留更多空间信息 model.fc = nn.Linear(512, num_classes) # 改输出维度 return model device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') model = build_resnet18().to(device)逻辑说明:ResNet-18 原版第一层是7×7, stride=2加maxpool,对 32×32 的图来说下采样太狠,改完以后特征图尺寸才够用。weights=None表示从随机初始化开始训,如果你数据量特别小可以考虑加载 ImageNet 预训练权重,但 CIFAR-10 有 5 万张训练图,从头训完全够。
参数上,bias=False是因为后面接了 BatchNorm,卷积层的 bias 会被抵消,省掉能减少参数量。nn.Identity()是 PyTorch 里常用的「占位不做事」模块,比写lambda x: x更规范。
3. 训练循环与调参:让 loss 真正降下去
3.1 损失函数、优化器与学习率调度
CIFAR-10 是单标签多分类,损失函数用CrossEntropyLoss没悬念。优化器我习惯用 SGD + momentum,虽然 Adam 收敛快,但在 CIFAR-10 上 SGD 最终精度通常更高。学习率调度用 CosineAnnealingLR,比 StepLR 更平滑。
import torch.optim as optim from torch.optim.lr_scheduler import CosineAnnealingLR criterion = nn.CrossEntropyLoss() optimizer = optim.SGD(model.parameters(), lr=0.1, momentum=0.9, weight_decay=5e-4) scheduler = CosineAnnealingLR(optimizer, T_max=200) # 200 个 epoch 内余弦衰减逻辑说明:weight_decay=5e-4就是 L2 正则化,在 PyTorch 里通过优化器的weight_decay参数实现,等价于在 loss 里加λ/2 * ||w||²。T_max=200要和你的总 epoch 数一致,否则学习率衰减节奏会乱。
参数上,初始lr=0.1是 SGD 在 CIFAR-10 上的经验值,如果你把 batch_size 翻倍到 256,学习率也可以翻到 0.2。momentum=0.9几乎是默认值,不用改。
3.2 训练循环:每个 epoch 该做什么
训练循环的骨架很固定,但细节决定成败。下面是一个带训练 + 验证的完整循环:
def train_one_epoch(model, loader, criterion, optimizer, device): model.train() total_loss, correct, total = 0.0, 0, 0 for imgs, labels in loader: imgs, labels = imgs.to(device), labels.to(device) optimizer.zero_grad() outputs = model(imgs) loss = criterion(outputs, labels) loss.backward() optimizer.step() total_loss += loss.item() * imgs.size(0) correct += (outputs.argmax(1) == labels).sum().item() total += imgs.size(0) return total_loss / total, correct / total def evaluate(model, loader, criterion, device): model.eval() total_loss, correct, total = 0.0, 0, 0 with torch.no_grad(): for imgs, labels in loader: imgs, labels = imgs.to(device), labels.to(device) outputs = model(imgs) loss = criterion(outputs, labels) total_loss += loss.item() * imgs.size(0) correct += (outputs.argmax(1) == labels).sum().item() total += imgs.size(0) return total_loss / total, correct / total for epoch in range(200): train_loss, train_acc = train_one_epoch(model, train_loader, criterion, optimizer, device) val_loss, val_acc = evaluate(model, test_loader, criterion, device) scheduler.step() if (epoch + 1) % 10 == 0: print(f"Epoch {epoch+1:3d} | train loss {train_loss:.4f} " f"acc {train_acc:.4f} | val loss {val_loss:.4f} acc {val_acc:.4f}")逻辑说明:model.train()和model.eval()必须成对出现,前者启用 BatchNorm 的 running stats 更新和 Dropout,后者冻结它们。optimizer.zero_grad()放在前向之前或之后都行,但一定要在每个 batch 开头清一次,否则梯度会累加。
参数上,loss.item() * imgs.size(0)是为了按样本数加权平均,避免最后一个 batch 不满时拉偏均值。torch.no_grad()在验证时关掉梯度计算,能省显存也能提速。
3.3 学习率与 batch size 的联动关系
很多人调参时只改一个不动另一个,结果要么 loss 震荡要么收敛慢。经验公式是:batch size 翻倍,学习率也翻倍。比如bs=128, lr=0.1对应bs=256, lr=0.2。但这不是线性的,当 batch size 超过 1024 以后,学习率再大反而会发散,这时候需要 warmup。
# 简单的 warmup:前 5 个 epoch 线性从 0.01 升到 0.1 def adjust_lr(optimizer, epoch, base_lr=0.1, warmup=5): if epoch < warmup: lr = base_lr * (epoch + 1) / warmup else: lr = base_lr for param_group in optimizer.param_groups: param_group['lr'] = lr逻辑说明:warmup 的作用是让模型在训练初期不要因为随机初始化的大梯度把权重带偏。param_groups是优化器里管理参数组的列表,直接改lr键就能生效。
4. 避坑与排查:那些让准确率卡在 60% 的坑
4.1 现象:训练 loss 不降,准确率在 10% 附近晃
原因:最常见的是标签和输出维度对不上,或者数据归一化把像素值压到了负数区间但模型第一层没适配。另一个可能是学习率太大,梯度直接炸了。
解决:先打印一个 batch 的imgs.min(), imgs.max(), labels[:10],确认数据范围在[-2.5, 2.5]左右、标签是 0~9 的整数。然后把学习率降到 0.01 试跑 2 个 epoch,如果 loss 开始降了,说明是 lr 问题。
4.2 现象:训练准确率 99%,测试准确率只有 70%
原因:过拟合。CIFAR-10 虽然简单,但如果你把数据增强关了、weight_decay 设成 0、模型又特别大,过拟合是必然的。
解决:把RandomCrop和RandomHorizontalFlip加回来,weight_decay调到5e-4或1e-3,如果还不行就加 Dropout 或者用更小的模型。我见过有人用 ResNet-50 跑 CIFAR-10,训练集 100% 测试集 75%,纯属杀鸡用牛刀还过拟合。
4.3 现象:CUDA out of memory但显存明明够
原因:PyTorch 的缓存分配器会预留显存,有时候前一个进程没退干净,或者num_workers太多导致每个 worker 都占一份内存。
解决:先nvidia-smi看有没有僵尸进程,有就kill -9。然后把batch_size减半,num_workers降到 2。如果还不行,在代码开头加torch.cuda.empty_cache(),但这只是缓解,根本办法是减小模型或 batch。
4.4 现象:每次跑出来的结果都不一样,没法复现
原因:随机种子没固定。PyTorch、NumPy、Python 内置 random 三处的种子都要设,而且cudnn的 benchmark 模式也会引入随机性。
解决:
import random, numpy as np, torch def set_seed(seed=42): random.seed(seed) np.random.seed(seed) torch.manual_seed(seed) torch.cuda.manual_seed_all(seed) torch.backends.cudnn.deterministic = True torch.backends.cudnn.benchmark = False set_seed(42)逻辑说明:deterministic=True会让 cuDNN 只用确定性算法,benchmark=False关掉自动调优。代价是训练速度可能慢 10%~20%,但换来的是可复现。
4.5 现象:验证集准确率比训练集还高
原因:如果验证集用了shuffle=True或者做了数据增强,评估结果会虚高。另一种可能是 BatchNorm 在eval()模式下用的 running stats 还没更新充分。
解决:确认test_loader的shuffle=False,test_transform里没有RandomCrop和RandomHorizontalFlip。如果是训练初期出现这种情况,等几个 epoch 让 running stats 稳定就好了。
5. 进阶技巧:把 CIFAR-10 准确率推到 95% 以上
5.1 测试时增强(TTA)与模型集成
单模型 ResNet-18 在 CIFAR-10 上大概能到 94%~95%,想再往上走,TTA 是最省事的办法。思路很简单:对同一张测试图做多次不同变换(原始、水平翻转、小幅度平移),分别推理后把 softmax 输出平均。
def tta_predict(model, imgs, device): model.eval() probs = [] with torch.no_grad(): # 原始 probs.append(torch.softmax(model(imgs.to(device)), dim=1)) # 水平翻转 probs.append(torch.softmax(model(torch.flip(imgs, dims=[3]).to(device)), dim=1)) # 上下翻转(CIFAR-10 里飞机和船可能受益,但汽车会受损,慎用) # probs.append(torch.softmax(model(torch.flip(imgs, dims=[2]).to(device)), dim=1)) return torch.stack(probs).mean(0)逻辑说明:torch.flip(imgs, dims=[3])是沿宽度方向翻转,对应水平镜像。TTA 的收益通常在 0.5%~1.5% 之间,取决于数据集和模型。注意不是所有翻转都有用,上下翻转对 CIFAR-10 里的汽车、卡车可能反而有害,因为车轮位置变了。
参数上,TTA 的推理时间会成倍增加,如果只做水平翻转就是 2 倍,加上原始就是 3 倍。线上服务要权衡延迟和精度。
5.2 混合精度训练:省显存、提速,但要注意 loss scaling
PyTorch 从 1.6 开始内置了torch.cuda.amp,用起来很简单:
from torch.cuda.amp import autocast, GradScaler scaler = GradScaler() for imgs, labels in train_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.step(optimizer) scaler.update()逻辑说明:autocast()会自动把部分算子转成 float16,GradScaler负责在反向传播时放大 loss 防止梯度下溢。scaler.update()会根据梯度情况动态调整缩放因子。
参数上,混合精度在支持 Tensor Core 的卡上(如 V100、A100、30/40 系)提速明显,能到 1.5~2 倍。但在老卡上可能没效果甚至更慢。另外,用了 AMP 以后weight_decay的行为可能略有变化,建议先跑一个 baseline 再对比。
5.3 用 TensorBoard 盯住训练过程
别只靠 print,TensorBoard 能让你看到 loss 曲线、学习率变化、权重分布。装好tensorboard后:
from torch.utils.tensorboard import SummaryWriter writer = SummaryWriter('runs/cifar10_resnet18') for epoch in range(200): train_loss, train_acc = train_one_epoch(...) val_loss, val_acc = evaluate(...) writer.add_scalar('Loss/train', train_loss, epoch) writer.add_scalar('Loss/val', val_loss, epoch) writer.add_scalar('Acc/train', train_acc, epoch) writer.add_scalar('Acc/val', val_acc, epoch) writer.add_scalar('LR', scheduler.get_last_lr()[0], epoch) writer.close()逻辑说明:SummaryWriter会把数据写到runs/目录,然后tensorboard --logdir=runs启动服务,浏览器打开就能看。scheduler.get_last_lr()返回的是列表,取第一个元素就是当前学习率。
我自己的习惯是:每次开新实验前先固定种子跑一遍 baseline,把 TensorBoard 曲线截图存下来,后面任何改动都跟这张图对比。如果 val loss 在某个 epoch 突然翘起来,八成是学习率太大或者数据增强太猛。从那以后我每次调参都强制走一遍「固定种子 → 跑 baseline → 改一个变量 → 对比曲线」的流程,再也没出现过「改了啥也不知道为啥变好」的情况。希望帮到你。
本文还有配套的精品资源,点击获取