☰
PyTorch实战:MNIST前馈神经网络训练与推理完整指南
2026/10/17 7:01:46 网站建设 项目流程

简介:这份资源面向深度学习入门者与需要快速验证手写数字识别效果的开发者,围绕MNIST数据集提供一套可直接运行的前馈神经网络训练方案,帮助读者跳过环境搭建与调参试错,直接获得可用的识别模型。压缩包共5个文件,以3个Python脚本和2个H5模型文件为主,脚本覆盖数据加载、网络定义与训练流程,H5文件分别保存模型参数与完整模型,整体约1.39MB,体积轻便便于本地部署与二次修改。目前已有6126人学习下载,说明其在实际练习与项目起步阶段具备较高参考价值。读者可借助完整代码理解前馈神经网络在MNIST上的训练细节,也可直接加载训练好的模型进行推理验证,适合课程作业、入门实验及快速原型开发等场景。

1. 拆开一个 MNIST 手写数字识别资源:为什么它值得你花 20 分钟跑一遍

如果你正在找一份能直接跑通的 MNIST 手写数字识别代码,并且希望拿到一个已经训练好的模型文件,省掉自己调参和等待收敛的时间,那这份资源就是冲着你来的。它包含完整的训练脚本、推理脚本,以及一个训练好的权重文件,下载后不需要额外配置数据集路径,脚本会自动拉取 MNIST 数据集并完成前馈神经网络的训练。适合刚接触深度学习、想先跑通一个完整流程的新手,也适合需要快速验证环境是否正常的老手。我见过太多人卡在“环境装好了但不知道下一步干什么”的阶段,这份东西就是用来打破这个僵局的。它不炫技,不堆复杂结构,就是一个干净的前馈网络,但恰恰因为简单,你才能把注意力放在数据流、损失曲线和模型保存这些真正重要的环节上。

2. 前馈神经网络跑 MNIST:从数据加载到模型保存的完整链路

2.1 为什么选前馈网络而不是 CNN 作为起点

MNIST 是一个 28x28 的灰度图数据集,共 70000 张样本,其中 60000 张用于训练,10000 张用于测试。很多人一上来就用卷积神经网络,觉得前馈网络太“浅”,但实际动手时你会发现,前馈网络在 MNIST 上达到 97% 以上的测试准确率只需要几行代码,训练时间在普通 CPU 上也就一两分钟。这个资源选的就是最朴素的全连接结构:输入层 784 个节点,隐藏层 128 个节点,输出层 10 个节点。隐藏层用 ReLU 激活,输出层用 LogSoftmax 配合 NLLLoss。这种组合在 PyTorch 里是标准写法,梯度稳定,不容易出现玄学问题。如果你连前馈网络的训练循环都没亲手写过,直接上 CNN 只会让你在池化步长和通道数上翻车。所以这个资源的定位很明确:先让你把“数据加载 → 前向传播 → 计算损失 → 反向传播 → 更新权重 → 保存模型”这条链路走通,再去考虑更复杂的结构。

2.2 训练脚本逐段拆解与参数说明

下面这段代码是资源里的训练脚本核心部分,我把它拆成三段来看。第一段是数据加载和预处理:

import torch from torchvision import datasets, transforms from torch.utils.data import DataLoader # 定义预处理:转为张量并归一化到 [0,1] 区间 transform = transforms.Compose([ transforms.ToTensor(), # 将 PIL 图像或 numpy 数组转为 torch.FloatTensor,同时除以 255 transforms.Normalize((0.1307,), (0.3081,)) # MNIST 的全局均值和标准差 ]) # 下载并加载训练集和测试集 train_dataset = datasets.MNIST(root='./data', train=True, download=True, transform=transform) test_dataset = datasets.MNIST(root='./data', train=False, download=True, transform=transform) # 创建数据加载器,batch_size 设为 64,训练集打乱顺序 train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True) test_loader = DataLoader(test_dataset, batch_size=1000, shuffle=False)

这里的Normalize用了 MNIST 的全局均值 0.1307 和标准差 0.3081,这是从训练集统计出来的固定值。如果你不做归一化,模型也能跑,但收敛会慢一些,而且损失曲线会抖得比较厉害。batch_size设 64 是一个经验值,显存占用小,梯度更新频率也够。测试集的batch_size设 1000 是为了一次性评估完所有测试样本,减少循环次数。注意download=True只在第一次运行时生效,之后会直接读取./data目录下的缓存文件。

第二段是模型定义:

import torch.nn as nn import torch.nn.functional as F class Net(nn.Module): def __init__(self): super(Net, self).__init__() self.fc1 = nn.Linear(784, 128) # 输入层到隐藏层 self.fc2 = nn.Linear(128, 10) # 隐藏层到输出层 def forward(self, x): x = x.view(-1, 784) # 将 28x28 的图像展平为 784 维向量 x = F.relu(self.fc1(x)) # 隐藏层使用 ReLU 激活 x = self.fc2(x) # 输出层不加激活,直接送 LogSoftmax return F.log_softmax(x, dim=1)

x.view(-1, 784)里的-1表示自动推断 batch 维度,这样写比x.view(x.size(0), -1)更简洁。F.log_softmax和nn.NLLLoss是配对使用的,如果你换成CrossEntropyLoss,那输出层就不要加log_softmax,否则会重复计算。这个细节很多新手会搞混,导致损失值异常。

第三段是训练循环和模型保存:

import torch.optim as optim device = torch.device("cuda" if torch.cuda.is_available() else "cpu") model = Net().to(device) optimizer = optim.Adam(model.parameters(), lr=0.001) for epoch in range(5): model.train() for batch_idx, (data, target) in enumerate(train_loader): data, target = data.to(device), target.to(device) optimizer.zero_grad() # 清空上一轮梯度 output = model(data) # 前向传播 loss = F.nll_loss(output, target) # 计算负对数似然损失 loss.backward() # 反向传播 optimizer.step() # 更新权重 if batch_idx % 100 == 0: print(f'Train Epoch: {epoch} [{batch_idx * len(data)}/{len(train_loader.dataset)}] Loss: {loss.item():.6f}') # 保存模型权重 torch.save(model.state_dict(), 'mnist_feedforward.pth') print('模型已保存为 mnist_feedforward.pth')

优化器用了 Adam,学习率 0.001。Adam 对初始学习率不敏感,适合这种小网络。如果你换成 SGD,学习率要调到 0.01 左右,并且加动量。训练 5 个 epoch 后,测试准确率通常能到 97.5% 以上。保存的是state_dict,只包含权重和偏置,不包含网络结构。加载时需要先实例化Net(),再调用load_state_dict。这个设计比保存整个模型更灵活,也避免了 pickle 在不同环境下的兼容问题。

2.3 推理脚本怎么写才能直接复用

资源里的推理脚本接收一张 28x28 的灰度图,输出预测的数字。核心代码如下:

import torch from PIL import Image import torchvision.transforms as transforms # 加载模型结构并恢复权重 model = Net() model.load_state_dict(torch.load('mnist_feedforward.pth', map_location='cpu')) model.eval() # 切换到评估模式,关闭 Dropout 和 BatchNorm 的训练行为 # 读取图片并预处理 img = Image.open('test_digit.png').convert('L') # 转为灰度图 transform = transforms.Compose([ transforms.Resize((28, 28)), transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,)) ]) img_tensor = transform(img).unsqueeze(0) # 增加 batch 维度 with torch.no_grad(): # 关闭梯度计算,节省内存 output = model(img_tensor) pred = output.argmax(dim=1, keepdim=True) print(f'预测数字: {pred.item()}')

model.eval()是必须的,虽然这个前馈网络没有 Dropout 和 BatchNorm,但养成习惯没坏处。unsqueeze(0)把单张图的维度从[1, 28, 28]变成[1, 1, 28, 28],因为模型期望的输入是[batch_size, 1, 28, 28]。torch.no_grad()在推理时一定要加,否则会构建计算图,浪费内存。如果你的图片是白底黑字,预测结果会正常;如果是黑底白字,需要先做反色处理,否则准确率会掉得很厉害。这个坑我在后面会细说。

3. 避坑与排查:MNIST 训练中五个高频翻车点

3.1 现象:损失值一直是 2.302 左右不下降

原因:输出层用了log_softmax,但损失函数用了CrossEntropyLoss,导致重复计算了 softmax。或者输入数据没有归一化,像素值在 0-255 之间,梯度爆炸后直接饱和。

解决:检查损失函数和输出层的配对关系。log_softmax配NLLLoss,CrossEntropyLoss配无激活的输出层。同时确认transforms.ToTensor()已经执行,它会把像素值缩放到 [0,1]。

3.2 现象:测试准确率只有 10% 左右,相当于随机猜

原因:训练时用了model.train(),但推理时忘了加model.eval(),或者加载权重时没有先实例化模型结构,导致load_state_dict报错后用了随机初始化的权重。

解决:推理前强制调用model.eval(),并且用try-except包裹load_state_dict,捕获RuntimeError后打印缺失或多余的键名。常见的是键名多了module.前缀,这是因为训练时用了DataParallel,保存的权重键名会带前缀。用model.load_state_dict({k.replace('module.', ''): v for k, v in state_dict.items()})可以去掉。

3.3 现象:训练到一半报错 “CUDA out of memory”

原因:batch_size设得太大,或者没有用with torch.no_grad()包裹测试循环,导致测试时也在构建计算图。

解决:把batch_size降到 32 或 64。测试循环里加with torch.no_grad():。如果显存实在小,可以在训练循环里用torch.cuda.empty_cache(),但这不是根本办法,降低 batch 更直接。

3.4 现象:自己手写的数字图片预测不准

原因:MNIST 数据集的图片是黑底白字,且数字居中、大小归一化到 20x20 像素区域再填充到 28x28。如果你用白底黑字的图片,或者数字没有居中,分布偏移会导致预测错误。

解决:对输入图片做反色处理,并裁剪出数字区域后居中缩放。可以用 PIL 的ImageOps.invert反色,再用ImageOps.expand填充到正方形后 resize。这个预处理步骤比模型本身更影响实际使用效果。

3.5 现象:训练集准确率很高,测试集准确率低很多

原因:过拟合。前馈网络只有两层,参数量不大,但如果训练 epoch 太多,或者训练集太小,也会过拟合。另外,如果测试集和训练集的预处理方式不一致,比如测试集忘了归一化,也会导致准确率骤降。

解决:控制 epoch 在 5 到 10 之间,观察测试集损失是否开始上升。确保训练和测试使用完全相同的transform。如果过拟合明显,可以在隐藏层后加nn.Dropout(0.2),但 MNIST 上通常不需要。

4. 把模型用起来:从单张推理到批量验证的进阶技巧

4.1 批量验证脚本与混淆矩阵输出

训练完模型后,不要只看一个准确率数字。我习惯跑一遍测试集,输出混淆矩阵,看看哪些数字容易被混淆。下面这段代码可以直接抄:

import torch from sklearn.metrics import confusion_matrix import seaborn as sns import matplotlib.pyplot as plt model.eval() all_preds = [] all_targets = [] with torch.no_grad(): for data, target in test_loader: data, target = data.to(device), target.to(device) output = model(data) preds = output.argmax(dim=1) all_preds.extend(preds.cpu().numpy()) all_targets.extend(target.cpu().numpy()) cm = confusion_matrix(all_targets, all_preds) plt.figure(figsize=(10, 8)) sns.heatmap(cm, annot=True, fmt='d', cmap='Blues') plt.xlabel('Predicted') plt.ylabel('True') plt.savefig('confusion_matrix.png') print('混淆矩阵已保存')

跑完你会看到,数字 4 和 9、3 和 8、5 和 6 之间的混淆比较多。这是 MNIST 的经典难点,不是你的模型有问题。知道这一点后,如果你要提升准确率,可以针对这些类别做数据增强,比如轻微旋转和平移。

4.2 用 TorchScript 导出模型,脱离 Python 环境运行

如果你想把模型部署到 C++ 环境或者移动端,state_dict不够用,需要导出 TorchScript。代码很简单:

model.eval() example_input = torch.rand(1, 1, 28, 28) traced_script_module = torch.jit.trace(model, example_input) traced_script_module.save('mnist_traced.pt') print('TorchScript 模型已保存')

torch.jit.trace会记录模型在前向传播时的操作,生成一个可以独立加载的.pt文件。加载时不需要再定义Net类,直接用torch.jit.load('mnist_traced.pt')即可。注意trace对控制流不友好,但这个前馈网络没有条件分支,所以没问题。导出后建议用同样的测试集验证一遍,确保精度没有损失。

4.3 一个我踩过的坑:模型文件版本兼容性

有一次我把在 PyTorch 1.9 下训练的state_dict拿到 PyTorch 1.13 环境里加载,报了一堆Unexpected key(s) in state_dict的错。原因是 PyTorch 在版本升级时改了nn.Linear的权重初始化方式,虽然键名没变,但某些层的默认参数变了。从那以后我每次保存模型都会在文件名里带上 PyTorch 版本号,比如mnist_feedforward_pytorch1.9.pth,并且在 README 里写清楚依赖版本。如果你拿到别人的模型文件加载失败,先检查版本,再检查键名。用print(model.state_dict().keys())和print(torch.load('xxx.pth').keys())对比一下,通常能发现问题。

4.4 验证模型是否真的“训练好了”的三个指标

不要只看准确率。我一般会同时看三个东西:测试集损失是否低于 0.1,混淆矩阵对角线是否占绝对优势,以及随机抽 20 张测试图做可视化预测,人眼过一遍。如果这三个都过了,模型才算真正可用。这个习惯帮我避免了好几次“准确率虚高但实际用起来一塌糊涂”的情况。希望帮到你。

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

需要专业的网站建设服务?

联系我们获取免费的网站建设咨询和方案报价,让我们帮助您实现业务目标

立即咨询