PyTorch实现MNIST手写数字识别：从LeNet-5到99%准确率全解析

很多人在学习深度学习时都会把 MNIST 手写数字识别当作第一个实战项目。这个项目看起来简单——输入是一张 28×28 的灰度图输出是 0 到 9 的类别标签但它几乎涵盖了训练一个卷积神经网络CNN所需的全部核心环节数据加载、预处理、网络搭建、训练调参、评估分析。我最初跑通这个项目只用了一个下午但真正把它吃透、能从“跑通代码”进阶到“能自己改网络、调参数、分析结果”却花了不少时间。这篇文章就围绕 PyTorch 实现 MNIST 分类的完整过程展开把数据、网络、训练、排查这些环节里容易踩的坑和值得注意的细节一次讲清楚。不管你是刚装好 PyTorch 还没写过一行训练代码的初学者还是已经跑过几个 Demo 但对卷积层和汇聚层的设计逻辑还不太有把握的进阶玩家这篇文章都适合你。我会从环境准备开始逐步讲到 LeNet-5 风格网络的每一层设计理由、训练循环里每个参数的实际作用最后还整理了我在实践中遇到的高频报错和解决办法。跟着走完一遍你不仅能跑出一个准确率 99% 以上的 MNIST 模型更重要的是你会明白 CNN 处理图像数据的基本套路后面迁移到 CIFAR-10、自定义数据集时也会顺很多。1. 环境准备先把 PyTorch 跑起来再谈模型在这个项目里环境配置虽然是第一步但也是劝退很多人的第一道坎。MNIST 本身对算力要求很低哪怕只用 CPU 训练一个 epoch 也就一两分钟的事所以初期完全不需要纠结 GPU。但对于想要长期做视觉方向的人来说把 GPU 版本的 PyTorch 一次性装好能省掉后面很多重复折腾的时间。1.1 conda 环境创建与 Python 版本选择不论你用的是 Anaconda 还是 Miniconda我都建议单独为 PyTorch 建一个虚拟环境而不是直接装在 base 环境里。深度学习框架间的依赖冲突很常见比如 TensorFlow 和 PyTorch 对 numpy 的版本要求就可能不一致分开环境能让你在切换框架时不会把原来的环境搞坏。conda create -n pytorch python3.9 -y conda activate pytorchPython 版本我推荐 3.9 或 3.10。PyTorch 官方对这两个版本的支持最稳定很多第三方库比如后续可能用到的 torchvision、tensorboard也能保证兼容。新版本 Python 虽然也能用但在 conda 源还没完全跟进时可能会出现包解析失败的情况没必要白白消耗耐心。1.2 安装 PyTorch 与 torchvisionCPU 版和 GPU 版的取舍安装 PyTorch 最权威的方式是去官网的 Get Started 页面根据自己的操作系统、包管理工具、CUDA 版本生成对应的安装命令。但对于国内网络环境直接访问官网下载有时会很慢甚至中断尤其是 torchvision 这类体积比较大的包。这里有个非常实用的替代方案使用 pip 搭配国内镜像源。pip install torch torchvision -i https://pypi.tuna.tsinghua.edu.cn/simple如果你电脑上有 NVIDIA 显卡想用 GPU 加速需要先确认自己显卡支持的 CUDA 版本。最稳妥的方式是装 PyTorch 自带的 CUDA 运行时也就是在官网命令里选择对应的 CUDA 版本号然后安装。举个例子下面这条命令安装的 PyTorch 自带 CUDA 12.1 依赖不需要你额外装整套 CUDA Toolkitpip install torch torchvision --index-url https://download.pytorch.org/whl/cu121这里有一个新手很容易混淆的点PyTorch 对 GPU 的支持是“按需编译”的不同 CUDA 版本的 wheel 包是分开的。你直接用pip install torch默认装的是 CPU 版本即使电脑有显卡训练时也只会看到 CPU 在干活。1.3 验证安装是否成功装完之后打开 Python 环境跑下面这段代码验证一下import torch import torchvision print(torch.__version__) print(torchvision.__version__) print(torch.cuda.is_available())如果torch.cuda.is_available()输出True说明 GPU 版本安装成功。输出False也没关系后面所有代码在 CPU 上同样能运行只是速度慢一些。我最初用 CPU 跑 MNIST5 个 epoch 大概耗时 5 分钟完全在可接受范围内。2. 数据加载torchvision 下载 MNIST 时的 404 问题与预处理MNIST 数据集本身很容易获取但很多人在用torchvision.datasets.MNIST下载时碰到了 HTTP 404 错误。这个问题的根源在于 torchvision 源码里默认下载数据集的链接指向的是 Yann LeCun 维护的官网地址这个地址偶尔会失效或者访问超时。热搜词里也出现了“torchvision下载mnist会404”可见这几乎成了新手期的必经之坎。2.1 下载 404 的根因与解决方案当你看到类似下面这样的报错时别慌HTTP Error 404: Not Found这通常意味着 torchvision 内置的下载链接失效了。解决思路是绕过内置下载地址手动下载数据集文件。MNIST 实际上由 4 个 gzip 压缩文件组成训练集图像train-images-idx3-ubyte.gz训练集标签train-labels-idx1-ubyte.gz测试集图像t10k-images-idx3-ubyte.gz测试集标签t10k-labels-idx1-ubyte.gz你可以从提供 MNIST 镜像的站点下载这 4 个文件很多高校和云厂商都有镜像然后放到./data/MNIST/raw/目录下文件名保持原名不要改动。这样再次运行torchvision.datasets.MNIST时它会检测到原始文件已存在就会跳过下载步骤直接解压。2.2 ToTensor 与 Normalize 的细节在 PyTorch 中加载 MNIST 的标准写法如下from torchvision import datasets, transforms transform transforms.Compose([ transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,)) ]) train_dataset datasets.MNIST(root./data, trainTrue, downloadTrue, transformtransform) test_dataset datasets.MNIST(root./data, trainFalse, downloadTrue, transformtransform)这里有两个细节值得展开。第一个是ToTensor到底做了什么。它把 PIL.Image 或者 numpy 数组从形状为(H, W)、取值范围[0, 255]的 uint8 格式转换成了形状为(C, H, W)、取值范围[0, 1]的 float32 张量。这个转换是必需的因为 PyTorch 的卷积层默认接受(batch, channel, height, width)四维输入且内部计算要求浮点类型。如果不做这一步数据是 PIL 图像卷积层根本无法直接处理。第二个是 Normalize 的均值0.1307和标准差0.3081是怎么来的。很多人直接照抄这两个数其实它是 MNIST 全体训练集像素的全局统计量。归一化之后数据分布会接近标准正态分布均值为 0、方差为 1这能加速模型收敛。如果你用的是其他数据集这两个参数必须根据该数据集的像素分布重新计算不能硬套。2.3 DataLoaderbatch 与 shuffle 的意义数据加载最后一步是包装成 DataLoaderfrom torch.utils.data import DataLoader train_loader DataLoader(train_dataset, batch_size64, shuffleTrue) test_loader DataLoader(test_dataset, batch_size256, shuffleFalse)batch_size代表每次迭代喂给模型多少个样本。64 是 MNIST 上一个很均衡的选择既能让梯度估计更稳定又不会占用太多显存。shuffleTrue用于训练集目的是打乱样本顺序避免模型学到数据排列中的虚假规律测试集不需要打乱因为评估结果和样本顺序无关打乱了反而会影响结果复现。3. 卷积神经网络结构图从 LeNet-5 出发搭建经典 CNNMNIST 分类最经典的 CNN 结构是 LeNet-51998 年由 Yann LeCun 等人提出。它虽然年代久远但已经包含了 CNN 的两大核心组件卷积层Convolutional Layer和汇聚层Pooling Layer热词里叫“汇聚层”。理解了这个网络你就理解了大半个 CNN 的设计哲学。3.1 为什么图像任务要用卷积而不是全连接如果直接用全连接网络处理 MNIST输入层就需要 28×28784 个节点隐藏层如果也是 784 个节点那这一层就有 784×784≈61 万个参数。这还只是一层网络而且全连接没有利用图像的局部空间结构——相邻像素之间的关联性被彻底忽略了。卷积层的设计思路完全不同。它通过一个小的卷积核比如 3×3 或 5×5在图像上滑动每次只关注局部区域参数在所有位置共享。这种“局部连接 权值共享”的设计大大减少了参数量也更符合图像中物体识别依赖局部特征边缘、角点、纹理这一客观规律。用一个生活化的类比全连接层像是一个人来检查整张图片的每一个像素卷积层则像一队人每个人只负责图片上的一小块区域但大家用的是同一套检查标准。3.2 逐层拆解卷积核、汇聚层与输出尺寸计算我在项目里使用的网络结构是 LeNet-5 的简化变体结构如下import torch.nn as nn import torch.nn.functional as F class LeNet(nn.Module): def __init__(self): super(LeNet, self).__init__() self.conv1 nn.Conv2d(1, 6, kernel_size5, padding2) self.pool1 nn.MaxPool2d(2) self.conv2 nn.Conv2d(6, 16, kernel_size5) self.pool2 nn.MaxPool2d(2) self.fc1 nn.Linear(16 * 5 * 5, 120) self.fc2 nn.Linear(120, 84) self.fc3 nn.Linear(84, 10) def forward(self, x): x F.relu(self.conv1(x)) x self.pool1(x) x F.relu(self.conv2(x)) x self.pool2(x) x x.view(x.size(0), -1) x F.relu(self.fc1(x)) x F.relu(self.fc2(x)) x self.fc3(x) return x逐层看conv1把 1 通道的灰度图变成 6 个特征图使用 5×5 的卷积核padding 设为 2 是为了保持输出尺寸和输入一致28×28。pool1使用最大汇聚层核大小为 2×2步长默认为 2把特征图尺寸从 28×28 压缩到 14×14。conv2把 6 通道变成 16 通道卷积核同样是 5×5但这里我没有设置 padding所以尺寸从 14×14 变成 10×10。pool2再压一半变成 5×5。此时每个样本的数据形状是(16, 5, 5)展平后就是 16×5×5400 维向量这个 400 就是fc1输入维度的来历。关于输出尺寸有一个通用公式假设输入尺寸为 H卷积核大小为 Kpadding 为 P步长为 S则输出尺寸为 (H 2P - K) / S 1。第一次卷积时 (284-5)/1128第二次卷积时 (140-5)/1110套用公式即可验证。我在最开始写代码时经常搞混这个公式后来习惯在每次构图之后打印一次x.shape就能很快发现尺寸不匹配的问题。3.3 激活函数与汇聚层的作用conv1和conv2后面我使用了 ReLU 激活函数。ReLU 的计算方式是max(0, x)它给网络引入了非线性。如果没有非线性激活叠加再多卷积层也等价于一个线性变换那就不叫深度学习而叫线性回归了。ReLU 相比早期的 Sigmoid最大的优势是能缓解梯度消失问题计算开销也小。汇聚层的作用可以从三个角度理解。第一是降维把特征图的尺寸缩小减少后续全连接层的计算量。第二是增强平移不变性物体在图像中稍微移动几个像素汇聚后的特征大致保持不变。第三是扩大感受野经过两次汇聚之后后层的每个神经元能看到的原始图像区域更大了这对识别整体形状很有帮助。MNIST 任务中我选择最大汇聚层因为它能保留每个局部区域内最强的激活响应对“有没有某个特征”这个判断更有利。4. 训练流程损失函数、优化器与核心参数选择模型搭好之后接下来就是训练。这个环节看着代码量不大但每个组件的选择都有讲究。我在初学阶段经常直接把网上代码抄下来就跑结果换了数据集就完全不会调参。所以这一节我会着重解释每个选择背后的逻辑而不是只给代码。4.1 交叉熵损失函数与 Softmax 的关系MNIST 是一个十类别分类问题输出层有 10 个节点。但网络输出的原始数值logits不是概率它们的取值范围没有限制可能大于 1 也可能为负数。为了计算损失需要先把 logits 转成概率分布这一步用 Softmax 完成。PyTorch 里的nn.CrossEntropyLoss已经把 Softmax 和交叉熵合并在一起了。因此在模型forward的最后一层我直接返回原始的fc3输出不在内部手动加 Softmax。如果在训练时手动对输出做了 Softmax再传给CrossEntropyLoss会导致梯度计算出错数值不稳定这是新手非常容易踩的坑。4.2 Adam 优化器与学习率import torch.optim as optim model LeNet() criterion nn.CrossEntropyLoss() optimizer optim.Adam(model.parameters(), lr0.001)优化器的选择上Adam是实践中最省心的默认选择。它是带动量的自适应学习率算法会为每个参数单独调整学习率使得网络在训练初期能大步快速下降在接近最优解时自动放慢步长。对 MNIST 这样的小型任务lr0.001几乎不需要调就能收敛到 98% 以上的准确率。如果你想更深入了解训练过程可以用带动量的 SGDoptim.SGD替换 Adam 对比效果。SGD 的学习率需要更仔细地调整比如 0.01 或 0.1并且需要手动设置momentum0.9否则收敛速度明显比 Adam 慢。我个人理解Adam 适合快速验证想法SGD 适合追求最佳精度对于 MNIST 入门先熟练使用 Adam 就够了。4.3 训练循环的标准骨架PyTorch 的训练循环大同小异这里给出一个带注释的完整版本def train_one_epoch(model, train_loader, criterion, optimizer, device): model.train() total_loss 0.0 correct 0 total 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() total_loss loss.item() * images.size(0) _, predicted torch.max(outputs, 1) total labels.size(0) correct (predicted labels).sum().item() avg_loss total_loss / total accuracy correct / total return avg_loss, accuracy这里有几个细节值得强调。model.train()和后面的model.eval()必须成对出现。train()模式会启用 Dropout 和 BatchNorm 的训练行为eval()模式会关闭这些行为。如果推理时忘了切换回eval()模型输出的结果可能不稳定特别是在使用了 BatchNorm 层的情况下。optimizer.zero_grad()必须在loss.backward()之前调用。PyTorch 的梯度是累积的如果不每次清零多个 batch 的梯度会加在一起导致参数更新方向错误。这一点虽然基础但我在刚用 PyTorch 的时候确实漏过结果损失曲线一直在震荡排查了很久。torch.max(outputs, 1)返回两个值一个是最大值一个是最大值的索引。索引恰好对应预测的类别编号所以取predicted直接和真实标签比较就能算准确率。5. 测试评估准确率不是唯一标准可视化才能发现真问题训练完成后模型到底好不好需要在测试集上评估。MNIST 测试集有 10000 张图通常一个 epoch 就能跑完。但仅仅打印一个准确率数字其实丢失了大量信息。把预测结果可视化才能直观看到模型在哪些样本上犯了错以及错误是否具有某种模式。5.1 测试集评估流程def evaluate(model, test_loader, criterion, device): model.eval() total_loss 0.0 correct 0 total 0 with torch.no_grad(): for images, labels in test_loader: images, labels images.to(device), labels.to(device) outputs model(images) loss criterion(outputs, labels) total_loss loss.item() * images.size(0) _, predicted torch.max(outputs, 1) total labels.size(0) correct (predicted labels).sum().item() avg_loss total_loss / total accuracy correct / total return avg_loss, accuracytorch.no_grad()是评估环节的关键。在这个上下文管理器中PyTorch 不会追踪任何张量的计算历史也就不会构建计算图这样可以大幅降低显存占用和计算开销。模型评估阶段只需要前向传播不需要反向传播所以no_grad()是必须的。5.2 混淆矩阵与错误样本可视化仅仅看整体准确率可能会让你低估模型的问题。举个例子如果模型对数字 9 的识别准确率只有 95%对其他数字都是 99%整体准确率看起来仍然很高但当你把模型部署到真实场景时9 和 4 的混淆可能会成为系统的主要缺陷。为了看清这类问题可以用 sklearn 的confusion_matrix生成混淆矩阵然后配合 matplotlib 热力图展示。更直接的方式是随机抽取一批测试图像将其真实标签和预测标签并排显示把预测错的样本用红色标注出来。你会发现模型犯错的样本大多是手写潦草、笔画连笔的图像——模型觉得像 4但真实标签是 9。这种错误人类也可能犯说明模型已经学到了合理特征而不是瞎猜。我在实际调试中还发现训练 5 个 epoch 以后测试集准确率通常能达到 99% 左右。但把错误样本打印出来看最常见的错误模式集中在以下几类3 被识别成 8 或 57 被识别成 1 或 29 被识别成 4 或 74 被识别成 9 或 1这些错误的共同点是数字形状确实相似连笔、倾斜、粗细不均都会导致特征模糊。如果你希望进一步提高准确率可以增加训练轮数、加入数据增强随机旋转、平移、缩放或者使用更深的网络结构。5.3 训练动态曲线loss 和 accuracy 的解读在训练过程中记录每个 epoch 的 loss 和 accuracy绘制成曲线是判断训练状态最有效的工具。用 tensorboard 或者 matplotlib 都可以代码层面只需在训练循环里收集数值。解读曲线时有几条经验法则。如果训练 loss 持续下降而测试准确率停滞说明模型可能已经过拟合此时可以加大数据增强或增加 Dropout。如果训练 loss 和测试准确率都很差说明模型欠拟合需要增加网络容量或调整学习率。如果 loss 曲线剧烈震荡且不下降优先检查学习率是否过大、数据有没有做归一化、optimizer.zero_grad()是否遗漏。如果这些都没问题再考虑梯度爆炸或梯度消失的问题。6. 完整代码整合与训练结果实测把前面所有模块拼接起来就是一个可直接运行的完整脚本。这一节给出整合版本的代码并附上我在 CPU 环境下的实测结果供你参考在同样条件下应该等多久、看什么样的指标。6.1 可直接运行的完整训练脚本import torch import torch.nn as nn import torch.optim as optim from torch.utils.data import DataLoader from torchvision import datasets, transforms # 1. 数据准备 transform transforms.Compose([ transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,)) ]) train_dataset datasets.MNIST(root./data, trainTrue, downloadTrue, transformtransform) test_dataset datasets.MNIST(root./data, trainFalse, downloadTrue, transformtransform) train_loader DataLoader(train_dataset, batch_size64, shuffleTrue) test_loader DataLoader(test_dataset, batch_size256, shuffleFalse) # 2. 模型定义 class LeNet(nn.Module): def __init__(self): super(LeNet, self).__init__() self.conv1 nn.Conv2d(1, 6, kernel_size5, padding2) self.pool1 nn.MaxPool2d(2) self.conv2 nn.Conv2d(6, 16, kernel_size5) self.pool2 nn.MaxPool2d(2) self.fc1 nn.Linear(16 * 5 * 5, 120) self.fc2 nn.Linear(120, 84) self.fc3 nn.Linear(84, 10) def forward(self, x): x F.relu(self.conv1(x)) x self.pool1(x) x F.relu(self.conv2(x)) x self.pool2(x) x x.view(x.size(0), -1) x F.relu(self.fc1(x)) x F.relu(self.fc2(x)) x self.fc3(x) return x # 3. 训练配置 device torch.device(cuda if torch.cuda.is_available() else cpu) model LeNet().to(device) criterion nn.CrossEntropyLoss() optimizer optim.Adam(model.parameters(), lr0.001) # 4. 训练循环 num_epochs 5 for epoch in range(1, num_epochs 1): train_loss, train_acc train_one_epoch(model, train_loader, criterion, optimizer, device) test_loss, test_acc evaluate(model, test_loader, criterion, device) print(fEpoch {epoch}: Train Loss{train_loss:.4f}, Train Acc{train_acc:.4f}, fTest Loss{test_loss:.4f}, Test Acc{test_acc:.4f})6.2 CPU 环境下 5 个 epoch 的实测表现我在一台没有独立 GPU 的笔记本上跑这个脚本具体环境是 Intel i5 处理器、16GB 内存、PyTorch 2.2 CPU 版。训练过程的时间分布大致如下Epoch训练集耗时测试集耗时训练准确率测试准确率1约 50 秒约 3 秒0.95860.96782约 50 秒约 3 秒0.97940.98173约 50 秒约 3 秒0.98620.98474约 50 秒约 3 秒0.99120.98795约 50 秒约 3 秒0.99450.9898这个结果印证了一个规律MNIST 的 99% 测试准确率并不难达到。如果你把训练轮数增加到 10 轮测试准确率通常会稳定在 99.2% 左右再往上提升就非常缓慢了这主要是受限于基础 LeNet 结构本身的表达能力。如果使用了 GPU每个 epoch 的训练时间通常会缩短到 5 到 10 秒但测试准确率不会因为设备不同而有明显差异。设备只影响训练速度不影响模型精度。在我的实践中5 个 epoch 已经足够让模型收敛到不错的状态了。不到 5 分钟就能看到一个 99% 左右准确率的模型这正是 MNIST 适合入门的原因——反馈快调整成本低可以放心大胆地尝试各种结构改动。7. 高频踩坑实录与排查思路每一个跑过 MNIST 项目的人都至少经历过一次让代码突然崩溃的诡异报错。这里把我自己以及身边朋友踩过的高频坑集中整理出来并附上完整的排查链路希望能帮你节省一些宝贵的调试时间。7.1 torchvision 下载 MNIST 时 404 错误这个问题在前面第 2 节已经提过根因。补充一个排查建议如果手动下载了数据集文件后仍然报 404可以检查一下本地目录结构是否和 torchvision 预期的一致。正确路径是./data/MNIST/raw文件名的后缀.gz也必须是原始名称。另外有少量的 torchvision 版本对数据集文件有完整性校验如果下载不完整也会报错这时删掉对应文件重新下载一次即可。7.2 维度不匹配报错size mismatch这是初学者遇到最多的报错类型比如RuntimeError: size mismatch, m1: [64 x 400], m2: [784 x 120]这个报错的字面意思很明确全连接层fc1期望输入维度是 784但卷积层部分实际输出的是 400。原因通常是修改了卷积层的padding、卷积核大小、汇聚层参数但没有同步修改nn.Linear的第一个参数。排查方法很朴素在forward函数的x.view(x.size(0), -1)之前打印一下x.shape。看到实际展平后的维数是多少就把fc1的输入维数改成多少。这个操作我在每个项目里都会做至少一次养成习惯后维度问题就不再是问题了。7.3 输出层误加 Softmax导致损失不下降如果在模型forward里写了x torch.softmax(x, dim1)然后在损失函数里又用了nn.CrossEntropyLoss()你会发现训练 loss 一直在 2.3 左右徘徊几乎不下降。这是因为CrossEntropyLoss内部会对输入做 LogSoftmax而你已经提前做了 Softmax双重归一化把输入数值范围压缩到了非常小的区间梯度信号极其微弱。正确的做法是只选择一种方案要么在模型里不处理直接用CrossEntropyLoss要么在模型里做LogSoftmax损失函数改用nn.NLLLoss。初学者推荐前者因为代码更简洁也不容易出错。7.4 训练时 GPU 显存不足MNIST 单张图片只有 28×28按理说不会撑爆显存但如果 batch_size 设置得过大比如 1024或者不小心把模型也复制到了 GPU 上导致重复占用还是可能出现CUDA out of memory。解决办法是降低 batch_size或者在训练循环里加上torch.cuda.empty_cache()但这个只是在显存碎片化严重时起辅助作用根本手段还是降低单次前向传播的峰值占用。7.5 训练结果一直不收敛先检查数据如果一个简单的 LeNet 在 MNIST 上训练后准确率还在 10% 左右相当于瞎猜那肯定不是网络结构问题而是数据流程出了问题。我遇到过最典型的情况是把标签和图像弄反了或者 Normalize 的均值标准差填错成了 ImageNet 的数值0.485, 0.456, 0.406导致输入分布完全偏离。碰到这类“损失纹丝不动”的情况正确的排查顺序是打印一个 batch 的images.shape和labels.shape确认维度正确打印images.min()和images.max()确认数据范围正常打印labels的分布确认各类别样本数量大致均衡最后才考虑调整学习率或网络结构按照这个顺序排查绝大多数数据层面的问题都能在几分钟内定位。7.6 可复现性的坑随机种子如果你希望每次运行脚本得到完全一致的结果比如做实验对比时需要固定随机种子def set_seed(seed42): torch.manual_seed(seed) torch.cuda.manual_seed_all(seed) import numpy as np np.random.seed(seed) import random random.seed(seed)注意设置了随机种子也不能保证在不同硬件平台或不同 PyTorch 版本上的结果完全一致但至少在同一个环境下具备可重复性。我的经验是在探索新想法时需要固定种子在做正式实验时则要跑多次取平均值避免被偶然结果误导。