ARTICLE · INTELLIGENCE

战地情报 · 详情页

来自尧图项目组的一线实战观察与深度解析

PyTorch实战:CIFAR-10图像识别从数据到部署全解析

PyTorch实战:CIFAR-10图像识别从数据到部署全解析 简介这份资源面向深度学习入门者与计算机视觉方向的初学者围绕PyTorch框架与CIFAR-10数据集提供一套可直接运行的图像识别实践材料帮助读者理解卷积神经网络从数据加载到模型训练、再到权重复用的完整链路。压缩包共5个文件约7.15MB以2个Python脚本为核心分别负责CIFAR-10图片数据的读取与预处理、CNN网络结构定义及训练流程实现另含1个已训练好的模型权重文件可直接加载用于预测或继续微调以及1份项目说明文档和1个数据集元信息文件便于快速了解目录结构与使用方法。目前已有53人学习。读者可借此掌握图像标准化、数据增强、损失函数与优化器配置等关键环节并借助现成权重跳过漫长训练适合作为课程实验、自学练手或进一步迁移学习的起点。1. 从一份 CIFAR-10 图像识别压缩包说起它到底能跑出什么很多人第一次接触图像识别都是从 CIFAR-10 开始的。它不像 ImageNet 那样动辄上百 GB也不像 MNIST 那样简单到几乎失去挑战性——10 个类别、6 万张 32×32 彩色图刚好卡在「能跑通」和「有难度」之间。一份名为「基于 PyTorch 的 CIFAR-10 图像识别.zip」的工程本质上就是把这条链路完整走了一遍数据加载、模型定义、训练循环、验证评估、模型保存。它解决的不是「识别猫狗」这种具体业务而是让你亲手摸到深度学习图像分类的每一颗螺丝。这篇文章适合两类人一类是刚装完 PyTorch、想找个能跑通的实战项目练手的新手另一类是做过 MNIST、想看看彩色图像和更深网络会带来什么变化的老手。我会按「数据怎么进、模型怎么搭、训练怎么稳、坑怎么避」的顺序把这份工程背后的技术点拆开讲清楚。你不需要有 GPU 也能跟但有一块 CUDA 显卡会舒服很多。2. 数据管道与预处理CIFAR-10 的 32×32 不是随便缩的2.1 为什么 CIFAR-10 的加载方式决定了训练上限CIFAR-10 的官方数据是 5 个训练 batch 加 1 个测试 batch每个 batch 是一个 10000×3073 的 uint8 数组。第一列是标签后面 3072 列是 32×32×3 的像素值按通道优先排列。PyTorch 的torchvision.datasets.CIFAR10已经帮你封装好了下载和解析但很多人直接downloadTrue就完事忽略了两个关键点一是数据增强的时机二是归一化参数的来源。常见做法是在DataLoader之前用transforms.Compose串起增强和归一化。训练集用随机裁剪加随机水平翻转测试集只做归一化。归一化参数一般用 CIFAR-10 的全局均值方差即mean(0.4914, 0.4822, 0.4465)、std(0.2023, 0.1994, 0.2010)。这个值不是拍脑袋来的是对训练集统计出来的。如果你用自己的数据集记得重新算一遍。import torch from torchvision import datasets, transforms # 训练集随机裁剪 翻转 张量转换 归一化 train_transform transforms.Compose([ transforms.RandomCrop(32, padding4), # 先四周补 4 像素再随机裁回 32 transforms.RandomHorizontalFlip(p0.5), # 一半概率水平翻转 transforms.ToTensor(), # 转成 [0,1] 的 float 张量 transforms.Normalize( # 按通道减均值除标准差 mean(0.4914, 0.4822, 0.4465), std(0.2023, 0.1994, 0.2010) ), ]) # 测试集只做张量转换和同样的归一化 test_transform transforms.Compose([ transforms.ToTensor(), transforms.Normalize( mean(0.4914, 0.4822, 0.4465), std(0.2023, 0.1994, 0.2010) ), ]) 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)是 CIFAR-10 的经典增强先补边再裁让模型看到物体在不同位置的样子。num_workers4在 Linux 上通常没问题Windows 下如果报错就改成 0 或加if __name__ __main__:保护。pin_memoryTrue配合 GPU 训练能减少 Host 到 Device 的拷贝开销但如果你只用 CPU这项可以关掉。2.2 数据加载的吞吐量怎么判断够不够一个简单的判断方法看第一个 epoch 的耗时里数据加载占了多少。如果 GPU 利用率长期低于 60%而 CPU 占用很高多半是num_workers不够或者磁盘 IO 慢。CIFAR-10 只有 170 MB 左右通常不是磁盘瓶颈更多是增强操作太复杂。我一般会先用num_workers4跑一个 epoch看nvidia-smi的 GPU 利用率曲线如果波动很大就加到 8 试试。但注意num_workers不是越大越好超过 CPU 物理核心数反而会因进程切换变慢。提示在 Windows 上使用多进程 DataLoader 时必须把训练代码放在if __name__ __main__:下面否则会无限递归创建子进程。3. 模型选型与搭建从 LeNet 到 ResNet 的中间地带3.1 为什么 CIFAR-10 上 VGG 和 ResNet 都能打但别一上来就上 ResNet-50CIFAR-10 的输入只有 32×32ImageNet 上那些为 224×224 设计的下采样策略在这里会过度压缩空间信息。ResNet-50 的第一层是 7×7 卷积加 stride 2再接 maxpool对 32×32 来说太快就把特征图降到 8×8 以下后面再堆残差块收益很低。常见做法是用为 CIFAR 定制的网络比如 VGG 风格的 6 层卷积加 3 层全连接或者 ResNet-20/32/44 这种浅层残差网络。我一般会先搭一个 VGG 风格的 baseline结构简单、训练稳定准确率能到 90% 左右。如果还想往上走再换 ResNet-20 加数据增强和余弦退火能到 93% 以上。下面是一个可直接用的 VGG 风格模型定义。import torch.nn as nn import torch.nn.functional as F class VGG_CIFAR(nn.Module): def __init__(self, num_classes10): super().__init__() # 特征提取3 - 64 - 128 - 256每次池化后空间减半 self.features nn.Sequential( nn.Conv2d(3, 64, kernel_size3, padding1), nn.BatchNorm2d(64), nn.ReLU(inplaceTrue), nn.Conv2d(64, 64, kernel_size3, padding1), nn.BatchNorm2d(64), nn.ReLU(inplaceTrue), nn.MaxPool2d(2), # 32 - 16 nn.Conv2d(64, 128, kernel_size3, padding1), nn.BatchNorm2d(128), nn.ReLU(inplaceTrue), nn.Conv2d(128, 128, kernel_size3, padding1), nn.BatchNorm2d(128), nn.ReLU(inplaceTrue), nn.MaxPool2d(2), # 16 - 8 nn.Conv2d(128, 256, kernel_size3, padding1), nn.BatchNorm2d(256), nn.ReLU(inplaceTrue), nn.Conv2d(256, 256, kernel_size3, padding1), nn.BatchNorm2d(256), nn.ReLU(inplaceTrue), nn.MaxPool2d(2), # 8 - 4 ) # 分类头4x4x256 - 512 - 10 self.classifier nn.Sequential( nn.Linear(256 * 4 * 4, 512), nn.ReLU(inplaceTrue), nn.Dropout(0.5), nn.Linear(512, num_classes), ) def forward(self, x): x self.features(x) x x.view(x.size(0), -1) # 展平 x self.classifier(x) return x这个模型的关键参数每个卷积后接 BatchNorm 和 ReLU池化用 2×2 最大池化。全连接前 Dropout 0.5 是 VGG 的标配能明显缓解过拟合。如果你把num_classes改成其他值这个结构也能直接迁移到类似尺寸的数据集上。3.2 参数量与显存占用的快速估算上面这个 VGG 风格模型大约 5.6M 参数用 batch_size128 训练时显存占用在 1.5 GB 左右GTX 1060 6G 完全够用。如果你换成 ResNet-20参数量降到 0.27M但训练时因为残差连接需要保存中间激活显存反而可能略高一点。估算显存有个经验公式参数量 × 4 字节 × 3梯度、动量、方差加上激活值。激活值跟 batch_size 和特征图大小成正比所以显存不够时优先降 batch_size而不是砍模型宽度。注意nn.CrossEntropyLoss内部已经包含 Softmax模型最后一层不要再加nn.Softmax否则会导致数值不稳定和收敛变慢。4. 训练循环与调参让 loss 曲线不再坐过山车4.1 优化器、学习率与权重衰减的搭配逻辑CIFAR-10 上最稳的组合是 SGD Momentum Weight Decay或者 AdamW。SGD 的初始学习率一般设 0.1动量 0.9权重衰减 5e-4。如果用了 BatchNorm权重衰减不要作用在 BN 的 gamma 和 beta 上否则会抑制 BN 的表达能力。PyTorch 里可以分组设置参数但简单起见很多人直接对整个模型加 weight_decay影响也不大。学习率调度用余弦退火CosineAnnealingLR比 StepLR 更平滑配合 warmup 效果更好。warmup 就是在最初几个 epoch 把学习率从很小线性升到初始值避免一开始梯度太大把 BN 的统计量带偏。下面是一个完整的训练循环骨架。import torch import torch.nn as nn import torch.optim as optim from torch.optim.lr_scheduler import CosineAnnealingLR device torch.device(cuda if torch.cuda.is_available() else cpu) model VGG_CIFAR(num_classes10).to(device) criterion nn.CrossEntropyLoss() optimizer optim.SGD(model.parameters(), lr0.1, momentum0.9, weight_decay5e-4) scheduler CosineAnnealingLR(optimizer, T_max200) # 200 个 epoch 内余弦下降 def train_one_epoch(epoch): model.train() running_loss 0.0 correct 0 total 0 for inputs, targets in train_loader: inputs, targets inputs.to(device), targets.to(device) optimizer.zero_grad() outputs model(inputs) loss criterion(outputs, targets) loss.backward() optimizer.step() running_loss loss.item() * inputs.size(0) _, predicted outputs.max(1) total targets.size(0) correct predicted.eq(targets).sum().item() scheduler.step() # 每个 epoch 后更新学习率 train_loss running_loss / total train_acc 100. * correct / total print(fEpoch {epoch}: loss{train_loss:.4f}, acc{train_acc:.2f}%) return train_loss, train_acc def evaluate(): model.eval() correct 0 total 0 with torch.no_grad(): for inputs, targets in test_loader: inputs, targets inputs.to(device), targets.to(device) outputs model(inputs) _, predicted outputs.max(1) total targets.size(0) correct predicted.eq(targets).sum().item() test_acc 100. * correct / total print(fTest acc: {test_acc:.2f}%) return test_acc for epoch in range(1, 201): train_one_epoch(epoch) if epoch % 10 0: evaluate()这里scheduler.step()放在 epoch 末尾如果是 PyTorch 1.1 之前的版本需要放在optimizer.step()之后、每个 batch 调用现在统一按 epoch 调用即可。T_max200表示余弦周期覆盖 200 个 epoch如果你只跑 100 个 epoch就改成 100。4.2 训练不收敛时先查这三处第一检查数据归一化是否和模型输入匹配。如果忘了归一化输入值在 0 到 255 之间第一层卷积的梯度会爆炸。第二检查学习率是否太大。SGD 用 0.1 在 CIFAR-10 上是安全的但如果你换了 Adam初始学习率要降到 1e-3 左右。第三检查 BatchNorm 的 momentum 和 eps。默认 momentum0.1、eps1e-5 通常没问题但如果 batch_size 小于 16BN 的统计量会很不稳定这时候要么增大 batch_size要么改用 GroupNorm。提示如果 loss 在前几个 epoch 就变成 nan先看输入里有没有 nan 或 inf再看学习率是不是太大导致梯度爆炸。可以在loss.backward()后加torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm5.0)做梯度裁剪。5. 避坑与排查CIFAR-10 工程里最容易翻车的五件事5.1 现象训练准确率一直卡在 10% 左右原因标签和输出对不上或者数据加载时标签被错误打乱。CIFAR-10 的targets是 0 到 9 的整数CrossEntropyLoss期望的也是这个范围。如果你自己写 Dataset 时把标签转成了 one-hot就会出问题。解决打印一个 batch 的targets看看是不是整数再检查模型输出维度是不是 10。5.2 现象验证集准确率比训练集高很多原因训练时用了 Dropout 和数据增强验证时关掉了所以验证集表现反而好。这在 CIFAR-10 上很常见尤其是训练早期。解决不用慌继续训练等训练集准确率也上来后验证集通常会略低于训练集。如果差距一直很大说明模型容量不够或者训练不够久。5.3 现象GPU 利用率忽高忽低训练速度慢原因DataLoader 的num_workers设置不当或者增强操作里有 Python 层面的慢操作。解决把num_workers设成 CPU 核心数的一半到全部用pin_memoryTrue并把transforms里耗时的操作比如自定义的旋转尽量用 torchvision 内置函数替代。5.4 现象保存的模型加载后准确率掉了一大截原因保存时只存了state_dict加载时模型结构变了或者忘了调用model.eval()。解决保存时用torch.save(model.state_dict(), model.pth)加载时先实例化同样的模型结构再model.load_state_dict(torch.load(model.pth))最后model.eval()。如果跨设备加载加map_locationdevice。5.5 现象多卡训练时 BatchNorm 表现异常原因nn.DataParallel会把 batch 切到多张卡上每张卡上的 batch_size 变小BN 统计量不准。解决改用nn.SyncBatchNorm或者DistributedDataParallel后者是官方推荐的多卡方案。如果只是单机多卡DistributedDataParallel配合torchrun启动更稳。6. 从 90% 到 94%几个我反复验证过的提分技巧第一个技巧是标签平滑Label Smoothing。把硬标签 0/1 换成 0.1/0.9 这种软标签能缓解过拟合在 CIFAR-10 上通常能提 0.5 到 1 个百分点。PyTorch 的CrossEntropyLoss自带label_smoothing参数直接设 0.1 就行。criterion nn.CrossEntropyLoss(label_smoothing0.1)第二个技巧是测试时增强TTA。对测试图做几次不同的裁剪或翻转把预测概率平均。CIFAR-10 上最简单的 TTA 是水平翻转一次准确率能提 0.3% 左右。代码上就是在evaluate里对每个 batch 跑两次前向一次原图一次翻转然后取平均。第三个技巧是混合精度训练AMP。用torch.cuda.amp把部分计算转成 float16显存占用减少约 40%训练速度提升 20% 到 30%准确率几乎不掉。下面是一个 AMP 训练的核心片段。from torch.cuda.amp import autocast, GradScaler scaler GradScaler() for inputs, targets in train_loader: inputs, targets inputs.to(device), targets.to(device) optimizer.zero_grad() with autocast(): outputs model(inputs) loss criterion(outputs, targets) scaler.scale(loss).backward() scaler.step(optimizer) scaler.update()注意autocast只包前向和 loss 计算scaler.scale(loss).backward()和scaler.step(optimizer)要放在外面。如果你在训练中遇到 loss 变成 inf先检查是不是某些操作在 float16 下溢出了可以把autocast的dtype改成torch.bfloat16试试。最后一个技巧是模型集成。把 VGG 和 ResNet 的预测概率平均准确率能到 94.5% 以上。但集成会增加推理成本适合对精度要求高、对延迟不敏感的场景。我一般会在单模型调到瓶颈后再考虑集成而不是一上来就堆模型。这些技巧里标签平滑和 AMP 是我每次都会开的TTA 看情况集成只在比赛或离线任务里用。训练 CIFAR-10 最大的教训是别急着换更深的网络先把数据增强、学习率调度和正则化调到位往往比换模型更有效。希望帮到你。本文还有配套的精品资源点击获取
RELATED READING

延伸阅读

更多一线实战笔记与深度复盘,助您持续精进