资讯动态

PyTorch实战:从零搭建CNN实现MNIST手写数字识别

发布时间:2026/10/16 3:47:36 来源:尧图企业网站定制
1. 项目缘起与整体设计思路1.1 为什么选择MNIST作为CNN入门的第一站搞深度学习这几年我越来越觉得入门阶段选对练手项目比看多少理论都管用。MNIST手写数字数据集就是这么一个存在——6万张训练图、1万张测试图每张都是28×28的灰度图10个类别。数据干净、体量适中、任务定义清晰几乎就是为CNN量身定做的教学素材。但很多人对MNIST有个误解觉得它太简单了跑个全连接网络都能到97%以上用CNN是不是杀鸡用牛刀我一开始也这么想直到自己动手搭了一遍才发现MNIST的价值根本不在于刷多高的准确率而在于它是一个完整的闭环验证平台。你可以在这个数据集上把卷积层、池化层、批归一化、Dropout、学习率调度、数据增强这些组件全部跑一遍观察每个组件对训练曲线和最终精度的影响。这种“控制变量式”的学习体验在复杂数据集上反而很难做到因为变量太多你分不清到底是哪个改动起了作用。所以这个项目的定位很明确用PyTorch从零搭建一个CNN在MNIST上完成训练、评估、可视化的全流程重点不是追求SOTA而是把每一层的设计意图、每一个超参数的来龙去脉讲清楚。适合刚学完PyTorch基础语法、想找个项目把知识串起来的同学也适合工作里需要快速验证某个CNN改动是否有效的老手拿来当模板。1.2 整体架构的取舍逻辑搭CNN之前先想清楚一件事你要的是一个“能跑就行”的demo还是一个“可以反复改”的实验框架我见过太多人上来就写一个nn.Sequential把层堆进去跑通了就完事结果想加个残差连接或者换个池化方式发现代码得推倒重来。我的做法是把网络定义、数据加载、训练循环、评估逻辑拆成独立的模块。网络部分用一个类封装forward里显式写出数据流向这样任何一层想改都只动一个地方。数据加载用DataLoader配合transforms把预处理和增强逻辑独立出来。训练循环写成函数接收模型、优化器、调度器作为参数方便做对比实验。具体到网络结构我选的是一个三层卷积两层全连接的经典配置。第一层用Conv2d(1, 32, 3, padding1)接ReLU再接MaxPool2d(2)第二层Conv2d(32, 64, 3, padding1)同样接激活和池化第三层Conv2d(64, 128, 3, padding1)之后做全局平均池化或者展平。全连接部分用Linear(128*3*3, 128)接Linear(128, 10)中间加Dropout。为什么是三层而不是两层或四层我实测过两层卷积在MNIST上大概能到99.1%左右三层能稳定到99.3%-99.4%四层反而因为参数增多、感受野过大在28×28的输入上出现了轻微的过拟合验证集精度波动变大。所以三层是一个性价比很高的选择。通道数从32到64到128逐层翻倍这是CNN设计的常规套路——浅层捕捉边缘纹理等低级特征需要较少通道深层组合成数字的局部结构需要更多通道来编码多样性。padding1配合kernel_size3这是保持特征图尺寸不变的经典组合。每次卷积后尺寸不变池化后减半28→14→7→3最后展平得到128×3×31152维特征。这个数字不是随便来的你可以自己算一遍输入28经过三次“卷积保持池化减半”28/21414/277/23向下取整所以是3×3。理解这个计算过程比死记网络结构重要得多。2. 核心细节解析与实操要点2.1 数据加载与预处理的关键参数MNIST在torchvision.datasets里可以直接下载但有几个坑得提前说清楚。第一个是transforms的顺序问题。很多人习惯先ToTensor()再Normalize()这个顺序是对的因为Normalize只接受Tensor输入。但均值和标准差填什么网上教程经常直接抄(0.1307, 0.3081)这是MNIST全局的均值和标准差用它是没问题的。但如果你想更严谨可以自己算一遍训练集的均值和方差代码就几行from torchvision import datasets, transforms import torch train_set datasets.MNIST(root./data, trainTrue, downloadTrue, transformtransforms.ToTensor()) loader torch.utils.data.DataLoader(train_set, batch_size1000, shuffleFalse) mean 0. std 0. total 0 for imgs, _ in loader: mean imgs.mean(dim[0,2,3]).sum() std imgs.std(dim[0,2,3]).sum() total imgs.size(0) print(mean/total, std/total)跑出来大概是0.1307和0.3081和网上一致。但自己算一遍的意义在于你以后换任何数据集都知道怎么处理而不是永远抄别人的数字。第二个坑是DataLoader的num_workers参数。在Windows上设成大于0有时候会报错Linux上设成4或8能明显加快数据加载。如果你发现GPU利用率忽高忽低大概率是数据加载成了瓶颈这时候把num_workers调大、开pin_memoryTrue会有改善。pin_memory的作用是把数据锁在页锁定内存里从CPU传到GPU时能走更快的通道这个细节很多教程不提但实测能省不少时间。第三个是数据增强。MNIST本身是居中的、规整的手写数字增强手段不宜太激进。我试过RandomRotation(10)和RandomAffine(degrees10, translate(0.1,0.1))前者能带来约0.1%的提升后者效果不明显反而拖慢训练。RandomHorizontalFlip千万别用数字6和9翻转后语义就变了这是常识但新手容易犯。我的建议是训练集加一个轻量的RandomRotation(10)就够了测试集不做任何增强只做ToTensor和Normalize。2.2 网络层设计的参数计算与选择依据卷积层的输出尺寸计算公式是O (I - K 2P) / S 1其中I是输入尺寸K是卷积核大小P是paddingS是步长。以第一层为例输入28K3P1S1代入得(28-32)/1128尺寸不变。池化层K2S2输出(28-2)/2114。这个计算过程建议每个搭CNN的人都手推一遍推完你就明白为什么padding1能让尺寸不变为什么池化核通常选2。通道数的选择有个经验法则从32起步每经过一次池化翻倍。这是因为池化让空间尺寸减半信息量压缩了需要通过增加通道数来补偿表达能力。但也不是越多越好我试过第一层就上64通道参数量翻倍训练时间增加约40%精度只提升了0.05%不到性价比很低。32→64→128这个配置在MNIST上属于“甜点区”。激活函数用ReLU是默认选择但有个细节nn.ReLU()和nn.ReLU(inplaceTrue)的区别。inplaceTrue会直接修改输入张量省一点内存在显存紧张时有用。但如果你在网络里用了残差连接inplaceTrue可能导致梯度计算出问题因为原始输入被覆盖了。MNIST这个简单网络用不用都行但我习惯显式写inplaceTrue养成省显存的习惯。Dropout加在哪里也有讲究。我见过有人在卷积层后面加Dropout这在MNIST上效果不好因为卷积层的参数共享本身就有正则化效果再加Dropout反而破坏了空间相关性。正确的做法是在全连接层加我一般在第一个全连接层后加Dropout(0.5)第二个全连接层前不加。实测这个配置比不加Dropout的验证集精度高约0.2%而且训练曲线更平滑。批归一化BatchNorm要不要加我的建议是加但位置要对。Conv2d→BatchNorm2d→ReLU→MaxPool2d这个顺序是经过验证的。BatchNorm放在卷积后、激活前能让激活值分布更稳定允许更大的学习率。我试过把BatchNorm放在激活后效果略差。加上BatchNorm后初始学习率可以从0.01提到0.1而不发散收敛速度明显加快。2.3 损失函数与优化器的搭配逻辑分类任务用CrossEntropyLoss是标准操作但要注意它内部已经包含了LogSoftmax所以网络最后一层不要加Softmax直接输出logits就行。我见过有人在最后一层加了Softmax再用CrossEntropyLoss相当于做了两次softmax梯度会变得很小训练几乎不动。这个坑很隐蔽因为代码不报错只是loss降不下去。优化器选Adam还是SGD这是个老生常谈的问题。Adam收敛快默认学习率0.001就能跑得很好适合快速验证想法。SGD配合动量需要调学习率但最终精度往往略高。我在MNIST上做过对比Adam跑20个epoch到99.3%SGD加动量跑30个epoch到99.4%差距很小但SGD多花了50%的时间。所以我的建议是入门用Adam追求极致精度且不差时间用SGD。学习率调度方面StepLR是最简单的每10个epoch乘以0.1。但我更推荐CosineAnnealingLR它让学习率按余弦曲线平滑下降不需要手动设milestones。实测在MNIST上CosineAnnealingLR比StepLR的最终精度高约0.1%而且训练后期loss震荡更小。ReduceLROnPlateau也可以但它需要监控验证集指标逻辑稍微复杂一点适合验证集波动较大的场景。权重衰减weight decay设多少Adam默认的1e-4在MNIST上偏大我一般设1e-5或者干脆不设因为Dropout和BatchNorm已经提供了足够的正则化。如果你发现训练集精度99.9%但验证集只有99.0%那就是过拟合了这时候把weight decay调到1e-4或者增大Dropout比例。3. 完整实操流程与核心代码实现3.1 环境准备与依赖安装先确认你的环境。PyTorch版本建议1.10以上torchvision对应版本即可。如果你有NVIDIA显卡装CUDA版本的PyTorch训练速度能快10倍以上。MNIST虽然小但用CPU训练一个epoch也要十几秒20个epoch就是好几分钟用GPU能压缩到几十秒。pip install torch torchvision验证安装是否成功import torch print(torch.__version__) print(torch.cuda.is_available())如果cuda.is_available()返回False但你确定有显卡大概率是PyTorch版本和CUDA驱动不匹配去官网重新选对应版本安装。这个环节卡住的人不少但解决起来就是重装的事。3.2 数据管道搭建import torch from torchvision import datasets, transforms from torch.utils.data import DataLoader train_transform transforms.Compose([ transforms.RandomRotation(10), transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,)) ]) test_transform transforms.Compose([ transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,)) ]) train_set datasets.MNIST(root./data, trainTrue, downloadTrue, transformtrain_transform) test_set datasets.MNIST(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)batch_size设128是个经验值。太小如32会导致梯度噪声大、训练慢太大如512会降低泛化能力而且MNIST只有6万张图512的batch一个epoch只有117次迭代收敛会变慢。128在速度和稳定性之间平衡得比较好。3.3 网络定义与参数初始化import torch.nn as nn import torch.nn.functional as F class MNISTCNN(nn.Module): def __init__(self, num_classes10, dropout0.5): super().__init__() self.conv1 nn.Conv2d(1, 32, 3, padding1) self.bn1 nn.BatchNorm2d(32) self.conv2 nn.Conv2d(32, 64, 3, padding1) self.bn2 nn.BatchNorm2d(64) self.conv3 nn.Conv2d(64, 128, 3, padding1) self.bn3 nn.BatchNorm2d(128) self.pool nn.MaxPool2d(2) self.fc1 nn.Linear(128 * 3 * 3, 128) self.fc2 nn.Linear(128, num_classes) self.dropout nn.Dropout(dropout) def forward(self, x): x self.pool(F.relu(self.bn1(self.conv1(x)))) x self.pool(F.relu(self.bn2(self.conv2(x)))) x self.pool(F.relu(self.bn3(self.conv3(x)))) x x.view(x.size(0), -1) x self.dropout(F.relu(self.fc1(x))) x self.fc2(x) return x参数初始化我习惯用Kaiming初始化虽然PyTorch默认的初始化在MNIST上也能跑但显式初始化能让训练更稳定def init_weights(m): if isinstance(m, nn.Conv2d) or isinstance(m, nn.Linear): nn.init.kaiming_normal_(m.weight, modefan_out, nonlinearityrelu) if m.bias is not None: nn.init.zeros_(m.bias) model MNISTCNN().to(device) model.apply(init_weights)Kaiming初始化的核心思想是让每一层的输出方差保持一致避免信号在传播过程中爆炸或消失。modefan_out表示用输出通道数来计算方差配合ReLU激活效果最好。这个细节在深层网络里影响很大MNIST上虽然不明显但养成习惯没坏处。3.4 训练循环与验证逻辑import torch.optim as optim from torch.optim.lr_scheduler import CosineAnnealingLR device torch.device(cuda if torch.cuda.is_available() else cpu) model MNISTCNN().to(device) criterion nn.CrossEntropyLoss() optimizer optim.Adam(model.parameters(), lr1e-3, weight_decay1e-5) scheduler CosineAnnealingLR(optimizer, T_max20) def train_epoch(model, loader, optimizer, criterion, device): model.train() total_loss 0 correct 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() return total_loss / len(loader.dataset), correct / len(loader.dataset) def evaluate(model, loader, criterion, device): model.eval() total_loss 0 correct 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() return total_loss / len(loader.dataset), correct / len(loader.dataset) for epoch in range(20): train_loss, train_acc train_epoch(model, train_loader, optimizer, criterion, device) val_loss, val_acc evaluate(model, test_loader, criterion, device) scheduler.step() print(fEpoch {epoch1:02d} | Train Loss {train_loss:.4f} Acc {train_acc:.4f} | fVal Loss {val_loss:.4f} Acc {val_acc:.4f})这里有几个细节值得说。model.train()和model.eval()必须成对出现前者启用Dropout和BatchNorm的训练模式后者切换到推理模式。忘了写model.eval()会导致验证时Dropout还在随机丢弃神经元BatchNorm还在用当前batch的统计量结果就是验证精度忽高忽低。torch.no_grad()在验证时一定要加它关闭梯度计算能省显存并加速。不加的话验证阶段也会构建计算图显存占用翻倍大模型上直接OOM。optimizer.zero_grad()的位置在PyTorch 1.7之后可以用optimizer.zero_grad(set_to_noneTrue)把梯度置为None而不是0能省一点内存。这个改动很小但体现了对细节的把控。3.5 训练过程监控与结果可视化跑完20个epoch你应该能看到训练集精度到99.8%以上验证集精度在99.3%-99.5%之间。如果验证集精度明显低于训练集比如训练99.9%验证99.0%那就是过拟合了需要增大Dropout或weight decay。如果两者都低比如训练95%验证94%那是欠拟合需要增加模型容量或延长训练。可视化方面我习惯画两条曲线loss曲线和accuracy曲线。用matplotlib几行代码就能搞定import matplotlib.pyplot as plt plt.figure(figsize(12, 4)) plt.subplot(1, 2, 1) plt.plot(train_losses, labelTrain Loss) plt.plot(val_losses, labelVal Loss) plt.xlabel(Epoch) plt.ylabel(Loss) plt.legend() plt.subplot(1, 2, 2) plt.plot(train_accs, labelTrain Acc) plt.plot(val_accs, labelVal Acc) plt.xlabel(Epoch) plt.ylabel(Accuracy) plt.legend() plt.show()健康的训练曲线应该是训练loss平滑下降验证loss先降后平两者差距不大。如果验证loss开始上升而训练loss还在降那就是过拟合的明确信号这时候可以早停early stopping保存验证loss最低的那个模型。还有一个进阶技巧把卷积层的feature map可视化出来看看网络到底学到了什么。第一层的卷积核通常呈现边缘检测器的形态第二层开始出现数字的局部笔画第三层能看到完整的数字轮廓。这个可视化对理解CNN的工作原理非常有帮助代码也不复杂取一层的输出用plt.imshow画出来就行。4. 常见问题与排查技巧实录4.1 训练不收敛的排查清单训练不收敛是新手最常遇到的问题表现是loss不降或者震荡剧烈。我整理了一个排查顺序按这个顺序走基本能定位到原因。现象可能原因排查方法解决方案loss完全不降学习率过大打印每步loss降到1e-4或1e-5loss降得很慢学习率过小观察前几个epoch提到1e-2或1e-3loss震荡剧烈batch_size太小检查batch_size增大到64或128loss变成NaN梯度爆炸打印梯度范数加梯度裁剪或降学习率精度卡在10%标签或输出维度错检查num_classes确认输出维度匹配精度卡在90%网络太浅或欠拟合增加层数或通道加深网络或延长训练梯度裁剪是个实用的技巧在loss.backward()之后、optimizer.step()之前加一行torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm1.0)这行代码把所有参数的梯度范数限制在1.0以内防止梯度爆炸。MNIST上一般用不到但养成习惯后换到深层网络能省很多调试时间。4.2 精度上不去的优化策略如果你的模型卡在99.0%左右上不去可以按以下顺序尝试优化。第一检查数据预处理是否正确Normalize的均值和标准差是否匹配。第二增加网络容量把通道数从32-64-128提到64-128-256但要注意过拟合风险。第三调整学习率调度从StepLR换成CosineAnnealingLR。第四加数据增强RandomRotation(10)是安全的选择。第五用模型集成训练3-5个不同初始化的模型预测时取平均能稳定提升0.1%-0.2%。我实测过一个配置三层卷积BNDropoutCosineAnnealingRandomRotation20个epoch能到99.45%。再往上提升就很困难了因为MNIST本身有约0.2%的标签噪声有些图连人都认不出来99.5%已经接近人类水平。所以不要盲目追求99.9%那大概率是过拟合了测试集。4.3 显存不足与训练速度优化MNIST模型很小显存一般不会不够。但如果你把batch_size设得很大或者用了很深的网络可能会遇到CUDA out of memory。解决方法按优先级第一减小batch_size。第二用torch.cuda.empty_cache()清理缓存。第三把模型改成半精度model.half()但要注意有些操作不支持半精度。第四用梯度累积小batch多次前向后累积梯度再更新等效于大batch。训练速度方面除了前面说的num_workers和pin_memory还可以用torch.backends.cudnn.benchmark True。这行代码让cuDNN自动寻找最适合当前硬件的卷积算法在输入尺寸固定的情况下能提速10%-20%。但如果你输入尺寸变化频繁反而会变慢因为每次都要重新搜索。还有一个容易忽略的点把数据预处理放在GPU上做。如果你的数据增强很复杂可以在DataLoader里只做ToTensor然后把Normalize和Rotation放到GPU上用torchvision.transforms的GPU版本或者自己写。不过MNIST的增强很简单这个优化收益不大换到ImageNet级别才明显。4.4 模型保存与加载的坑保存模型有两种方式保存整个模型和只保存参数。我强烈建议只保存state_dicttorch.save(model.state_dict(), mnist_cnn.pth)加载时先实例化模型结构再加载参数model MNISTCNN().to(device) model.load_state_dict(torch.load(mnist_cnn.pth, map_locationdevice)) model.eval()保存整个模型的问题是pickle会绑定具体的类路径换个项目目录或者改了类名就加载失败。只保存参数则没有这个问题只要网络结构一致就能加载。map_location参数在跨设备加载时很有用比如在GPU上训练、在CPU上推理不加这个参数会报错。还有一个细节加载参数后一定要调model.eval()否则Dropout和BatchNorm还是训练模式推理结果会不稳定。这个坑我踩过当时推理结果每次都不一样排查了半天才发现是忘了切eval模式。5. 从MNIST到真实项目的迁移经验MNIST跑通只是起点真正有价值的是把这套流程迁移到更复杂的任务上。我说几个迁移时需要注意的点。第一输入尺寸变了全连接层的输入维度要重新算。比如CIFAR-10是32×32的RGB图通道数从1变3最后展平的维度完全不同。第二类别数变了最后一层的输出维度要改。第三数据增强策略要重新设计MNIST的Rotation在自然图像上可能不适用。第四学习率和batch_size需要重新调大数据集通常用更大的batch_size和更小的学习率。但核心骨架是不变的数据管道、网络定义、训练循环、验证逻辑、模型保存。这套模板我用了好几年换数据集时只改配置不改结构。建议你也把这套代码整理成一个自己的模板以后新项目直接复制粘贴能省很多时间。另外MNIST上验证过的技巧比如BatchNorm的位置、Dropout的比例、学习率调度的选择在更复杂的任务上依然适用只是具体数值需要微调。所以不要觉得MNIST简单就跳过把它吃透后面学ResNet、Transformer都会轻松很多。我自己在带新人的时候都是让他们先把MNIST的CNN手写三遍第一遍照着教程抄第二遍不看教程自己写第三遍改结构做对比实验。三遍下来CNN的基本功就扎实了。

读完文章,也想定制专属网站?

尧图设计师 24 小时内与您沟通定制方案

免费获取报价 →
↑