PyTorch CIFAR-10图像识别实战：从环境搭建到准确率提升

简介：这份资源面向深度学习入门者与计算机视觉方向的初学者，围绕PyTorch框架与CIFAR-10数据集，提供一套可直接运行的图像识别实践材料，帮助读者理解卷积神经网络从数据加载到模型训练、再到权重复用的完整链路。压缩包共5个文件，以2个Python脚本为核心，分别负责CIFAR-10图片数据的读取与预处理、CNN网络结构定义及训练流程实现，另含1个已训练好的模型权重文件、1个数据集元信息文件和1份项目说明文档，整体约7.15MB，体积轻便，便于本地快速跑通。其中数据预处理环节涉及标准化与数据增强，训练脚本则覆盖损失函数、优化器选择等关键配置，已训练权重可直接用于预测或微调，省去从零训练的时间成本。目前已有53人学习，适合希望以最小成本上手图像识别任务、并借此理解机器学习项目组织方式的读者参考。

1. 从一份 CIFAR-10 识别代码包说起：它到底能跑出什么结果

很多人第一次接触图像识别，都是从 CIFAR-10 这个数据集开始的。它小，10 个类别、6 万张 32×32 的彩色图，单卡十几分钟就能跑完一轮，特别适合拿来验证 PyTorch 环境搭得对不对、训练流程写得顺不顺。但真正上手时你会发现，网上大量所谓「PyTorch 图像识别实战」的代码包，要么环境依赖写得含糊，要么训练脚本和推理脚本混在一起，要么准确率卡在 60% 上不去，你根本不知道是模型问题还是数据问题。这份基于 PyTorch 的 CIFAR-10 图像识别代码包，要解决的就是这条从环境搭建到模型跑通、再到准确率能看的完整链路。它适合两类人：一类是刚装完 PyTorch、想找个能跑通的图像识别项目练手的新手；另一类是做过 MNIST、想往彩色图像分类迈一步的熟手。下面我按自己复现这类项目的顺序，把每一步拆开讲。

2. 环境搭建与数据准备：PyTorch 装完只是开始

2.1 先确认 CUDA 和 PyTorch 版本对得上

这一步翻车的人最多。你兴冲冲pip install torch，装完发现torch.cuda.is_available()返回 False，然后开始怀疑显卡坏了。其实大概率是装成了 CPU 版本，或者 CUDA 版本和驱动不匹配。我一般会先看驱动支持的 CUDA 上限，再去 PyTorch 官网找对应的安装命令，而不是随手pip install。

# 先看驱动和 CUDA 版本 nvidia-smi # 输出里右上角 CUDA Version 是驱动支持的最高版本 # 比如显示 12.4，那你可以装 cu121 或 cu118 的 PyTorch

确认之后，用官方推荐的命令装，别用默认源：

# 以 CUDA 12.1 为例，具体命令以 PyTorch 官网为准 pip install torch torchvision --index-url https://download.pytorch.org/whl/cu121

装完立刻验证，这一步不能省：

import torch print(torch.__version__) # 看版本号 print(torch.cuda.is_available()) # 必须是 True print(torch.cuda.get_device_name(0)) # 看显卡型号对不对

逻辑说明：torch.cuda.is_available()是判断 GPU 能不能用的唯一标准，返回 False 就别往下跑了，先解决环境。参数上，--index-url指向的是 PyTorch 官方 wheel 源，国内网络下如果慢，可以换清华源，但要注意清华源同步的版本可能滞后，装之前先确认有没有你要的 CUDA 版本。如果你用的是 WSL 或者 Ubuntu，装完还要确认nvidia-smi在 WSL 里能正常输出，否则 PyTorch 一样看不到显卡。

2.2 CIFAR-10 数据集的下载与预处理

CIFAR-10 用torchvision.datasets一行就能下载，但有几个参数必须设对，否则后面训练会出各种玄学问题。

import torch from torchvision import datasets, transforms # 训练集预处理：随机裁剪 + 翻转 + 归一化 train_transform = transforms.Compose([ transforms.RandomCrop(32, padding=4), # 先 padding 再随机裁回 32 transforms.RandomHorizontalFlip(), # 50% 概率水平翻转 transforms.ToTensor(), # 转成 tensor，像素值归到 [0,1] transforms.Normalize((0.4914, 0.4822, 0.4465), # CIFAR-10 均值 (0.2470, 0.2435, 0.2616)) # CIFAR-10 标准差 ]) # 测试集只做 ToTensor 和 Normalize，不做增强 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', train=True, download=True, transform=train_transform) test_set = datasets.CIFAR10(root='./data', train=False, download=True, transform=test_transform) train_loader = torch.utils.data.DataLoader(train_set, batch_size=128, shuffle=True, num_workers=4) test_loader = torch.utils.data.DataLoader(test_set, batch_size=128, shuffle=False, num_workers=4)

逻辑说明：RandomCrop(32, padding=4)是先四周各补 4 像素再随机裁回 32×32，这是 CIFAR-10 上最常用的数据增强，能明显缓解过拟合。Normalize里的均值和标准差是 CIFAR-10 训练集统计出来的固定值，直接用就行，别自己算。num_workers在 Windows 上建议设 0，否则容易报多进程相关的错；Linux 和 WSL 下设 4 或 8 都行，看 CPU 核数。batch_size设 128 是常见起点，显存不够就降到 64。

提示：第一次运行会下载约 170MB 的数据到./data目录，下载慢的话可以手动下载后放到对应目录，download=True检测到文件存在就不会重复下。

3. 模型选型与训练脚本：从 LeNet 到 ResNet 的取舍

3.1 小模型够不够用，什么时候该上 ResNet

CIFAR-10 上，一个简单的卷积网络就能跑到 70% 以上，但想上 90% 就得换结构。我一般会准备两个模型：一个轻量的自定义 CNN 用来快速验证流程，一个 ResNet-18 用来冲准确率。轻量模型的好处是训练快，改起来直观，适合调试数据管道；ResNet-18 有残差连接，深层也能训得动，是 CIFAR-10 上的经典选择。

import torch.nn as nn import torch.nn.functional as F class SimpleCNN(nn.Module): def __init__(self, num_classes=10): super().__init__() # 3 -> 32 -> 64 -> 128 通道逐层加倍 self.conv1 = nn.Conv2d(3, 32, 3, padding=1) self.conv2 = nn.Conv2d(32, 64, 3, padding=1) self.conv3 = nn.Conv2d(64, 128, 3, padding=1) self.pool = nn.MaxPool2d(2, 2) self.fc1 = nn.Linear(128 * 4 * 4, 256) self.fc2 = nn.Linear(256, num_classes) self.dropout = nn.Dropout(0.5) def forward(self, x): x = self.pool(F.relu(self.conv1(x))) # 32x32 -> 16x16 x = self.pool(F.relu(self.conv2(x))) # 16x16 -> 8x8 x = self.pool(F.relu(self.conv3(x))) # 8x8 -> 4x4 x = x.view(x.size(0), -1) # 展平 x = self.dropout(F.relu(self.fc1(x))) x = self.fc2(x) return x

逻辑说明：三次卷积每次后接 2×2 最大池化，特征图从 32 降到 4，通道从 3 升到 128。view(x.size(0), -1)把每个样本展平成一维，x.size(0)是 batch 维度，不能写死。Dropout(0.5)放在全连接层前，是防过拟合的常规操作。这个模型参数量小，单卡跑 50 个 epoch 也就几分钟。

如果你要上 ResNet-18，torchvision.models里直接有，但 CIFAR-10 的图只有 32×32，原版 ResNet 第一层是 7×7 卷积加 stride 2，会把图降得太快，需要改一下：

import torchvision.models as models def get_resnet18(): model = models.resnet18(num_classes=10) # 把第一层 7x7 stride2 换成 3x3 stride1，适配 32x32 输入 model.conv1 = nn.Conv2d(3, 64, kernel_size=3, stride=1, padding=1, bias=False) # 去掉原来的 maxpool 层 model.maxpool = nn.Identity() return model

逻辑说明：nn.Identity()是恒等映射，相当于把 maxpool 层「跳过」。这样改完，ResNet-18 在 CIFAR-10 上训练更稳，准确率也更高。参数上，num_classes=10对应 CIFAR-10 的类别数，别写成 1000。

3.2 训练循环里必须盯住的几个参数

训练脚本本身不复杂，但有几个参数设错，结果会差很多。我一般把学习率、优化器、学习率调度这三样固定成一套组合，跑通了再调。

import torch.optim as optim from torch.optim.lr_scheduler import CosineAnnealingLR device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') model = SimpleCNN().to(device) criterion = nn.CrossEntropyLoss() optimizer = optim.SGD(model.parameters(), lr=0.1, momentum=0.9, weight_decay=5e-4) scheduler = CosineAnnealingLR(optimizer, T_max=50) for epoch in range(50): model.train() running_loss = 0.0 for images, labels in train_loader: images, labels = images.to(device), labels.to(device) optimizer.zero_grad() outputs = model(images) loss = criterion(outputs, labels) loss.backward() optimizer.step() running_loss += loss.item() scheduler.step() print(f'Epoch {epoch+1}, Loss: {running_loss/len(train_loader):.4f}')

逻辑说明：SGD配momentum=0.9和weight_decay=5e-4是 CIFAR-10 上的经典组合，weight_decay就是 L2 正则化，防止权重过大。CosineAnnealingLR让学习率按余弦曲线从 0.1 降到接近 0，T_max设成总 epoch 数。optimizer.zero_grad()必须在loss.backward()之前，否则梯度会累加。images.to(device)和labels.to(device)不能漏，漏了会报设备不一致的错。

注意：如果你用 ResNet-18，学习率可以设 0.1 起步；如果用 SimpleCNN，0.01 到 0.05 更稳。学习率太大 loss 会震荡甚至变 NaN，太小则收敛慢。

4. 评估、保存与推理：别让模型只活在训练脚本里

4.1 测试集评估的正确写法

训练完不评估，等于白跑。评估时有两个坑：忘了切eval()模式，以及忘了torch.no_grad()。

def evaluate(model, loader, device): model.eval() # 切到评估模式，关掉 dropout correct = 0 total = 0 with torch.no_grad(): # 不计算梯度，省显存 for images, labels in loader: images, labels = images.to(device), labels.to(device) outputs = model(images) _, predicted = torch.max(outputs, 1) # 取最大 logit 的索引 total += labels.size(0) correct += (predicted == labels).sum().item() return correct / total acc = evaluate(model, test_loader, device) print(f'Test Accuracy: {acc:.4f}')

逻辑说明：model.eval()会关掉 Dropout 和 BatchNorm 的训练行为，不切的话结果会偏低且不稳定。torch.no_grad()关闭梯度计算，评估阶段不需要反向传播，能省不少显存。torch.max(outputs, 1)返回每行最大值和对应索引，索引就是预测类别。labels.size(0)是当前 batch 的样本数，累加得到总数。

4.2 保存模型和单张图片推理

模型保存有两种方式：保存整个模型和只保存参数。我推荐只保存state_dict，加载时更灵活。

# 保存 torch.save(model.state_dict(), 'cifar10_model.pth') # 加载 model = SimpleCNN().to(device) model.load_state_dict(torch.load('cifar10_model.pth', map_location=device)) model.eval()

单张图片推理时，预处理必须和测试集一致，否则结果会离谱：

from PIL import Image def predict_image(image_path, model, device): transform = transforms.Compose([ transforms.Resize((32, 32)), transforms.ToTensor(), transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2470, 0.2435, 0.2616)) ]) img = Image.open(image_path).convert('RGB') img_tensor = transform(img).unsqueeze(0).to(device) # 加 batch 维度 with torch.no_grad(): output = model(img_tensor) _, predicted = torch.max(output, 1) return predicted.item()

逻辑说明：unsqueeze(0)是在第 0 维加一个 batch 维度，因为模型输入要求是(N, C, H, W)。convert('RGB')防止 PNG 带透明通道导致通道数不对。预处理里的均值和标准差必须和训练时完全一致，差一点准确率就掉。

提示：如果你要把模型转到 ONNX 部署，torch.onnx.export时记得设opset_version，并且输入用torch.randn(1, 3, 32, 32)做 dummy input，动态 batch 维度用dynamic_axes指定。

5. 避坑与排查：那些让准确率卡在 60% 的常见问题

5.1 训练 loss 不降，先查这四处

现象：跑了几个 epoch，loss 一直在 2.3 附近晃，准确率跟随机猜差不多。原因通常是四选一：学习率太大导致震荡、数据标签没对齐、模型输出维度不对、或者输入没归一化。解决顺序是先把学习率降到 0.001 试一个 epoch，如果 loss 开始降，就是学习率问题；如果还不降，打印一个 batch 的images.shape和labels，确认是(128, 3, 32, 32)和 0 到 9 的整数；再检查fc2输出是不是 10。归一化漏了的话，loss 也会降得很慢。

5.2 训练准确率高但测试准确率低

现象：训练集准确率 95%，测试集只有 70%。这是典型过拟合。原因可能是数据增强太弱、模型太大、或者训练太久。解决：先加RandomCrop和RandomHorizontalFlip，再把weight_decay从 5e-4 提到 1e-3，还不行就减小模型或加 Dropout。CIFAR-10 上 ResNet-18 不加增强也能到 90% 以上，但 SimpleCNN 不加增强很容易过拟合。

5.3 CUDA out of memory 的三种处理

现象：训练到一半报RuntimeError: CUDA out of memory。原因就一个：显存不够。解决按优先级来：先把batch_size减半，这是最直接的；然后检查有没有在循环里累积 tensor 没释放，比如把 loss 存进列表还带着计算图，改成loss.item()；最后确认torch.no_grad()在评估时加了。如果都不行，换小模型或者用梯度累积模拟大 batch。

5.4 num_workers 在 Windows 上报错

现象：Windows 下设num_workers=4，报BrokenPipeError或卡死。原因是 Windows 的多进程和 Linux 不一样，DataLoader 的 worker 启动方式有差异。解决：把num_workers设成 0，数据加载在主进程里做，慢一点但稳。或者在if __name__ == '__main__':里包住训练代码，这是 Windows 多进程的硬性要求。

5.5 准确率上不去，别急着换模型

现象：SimpleCNN 跑完 50 epoch，测试准确率 75%，你觉得是模型不行。先别换 ResNet，按这个顺序查：学习率调度有没有加、数据增强有没有开、batch_size 是不是太小、epoch 是不是不够。我见过太多人直接上 ResNet，结果因为学习率没调好，还不如调好的 SimpleCNN。先把现有模型的每个参数调一遍，再考虑换结构。

6. 把准确率从 75% 推到 93% 的几个具体手法

如果你已经跑通了上面的流程，SimpleCNN 大概在 75% 到 80% 之间，ResNet-18 能到 90% 左右。想再往上推，我一般用下面这几招，按性价比排序。

第一招是换学习率调度。CosineAnnealingLR比固定学习率好，但OneCycleLR在 CIFAR-10 上收敛更快。用法是把total_steps设成epochs * len(train_loader)，max_lr设成 0.1，然后每个 batch 后调scheduler.step()，不是每个 epoch。

from torch.optim.lr_scheduler import OneCycleLR scheduler = OneCycleLR(optimizer, max_lr=0.1, total_steps=50 * len(train_loader), pct_start=0.3) for epoch in range(50): for images, labels in train_loader: # ... 训练步骤 ... optimizer.step() scheduler.step() # 注意：每个 batch 后调

第二招是加 CutMix 或 MixUp 增强。这两种增强把两张图按比例混合，标签也按比例混合，能显著提升泛化。CutMix 的实现大概十几行，核心是随机生成一个矩形框，把一张图的部分区域替换成另一张图的对应区域，loss 按面积比例加权。

第三招是测试时增强（TTA）。推理时把原图、水平翻转图各跑一遍，输出概率平均，准确率能涨 0.5 到 1 个点。代价是推理时间翻倍，看场景取舍。

第四招是模型集成。训 3 到 5 个不同初始化的 ResNet-18，推理时把 logits 平均，准确率能到 94% 以上。这是冲榜的常规操作，但部署时模型体积也翻倍。

最后说个验证方法：别只看最终准确率，把每个 epoch 的测试准确率画出来，看曲线有没有过拟合的拐点。如果测试准确率在第 30 个 epoch 开始降，说明 30 之后都在过拟合，早停就行。我自己的习惯是每跑完一组参数，把配置和结果记在一个表格里，下次调参直接翻记录，比凭感觉试快得多。这套流程我从 MNIST 一路用到 CIFAR-10，再到更大的数据集，基本没变过，希望帮到你。

本文还有配套的精品资源，点击获取