简介:面向深度学习与机器学习入门者的MNIST手写数字识别实践素材,围绕前馈神经网络完成数据加载、模型构建、训练与推理的完整闭环,适合课程实验、项目复现和快速验证。压缩包共5个文件,以3个Python脚本和2个H5模型文件为主,前者承载训练与推理代码,后者保存训练好的权重及完整模型,整体约1.39MB,轻量易取用。已有6131人学习下载,说明其在入门实践场景中具有较高参考价值。读者可直接获得可运行的代码与训练成品,省去从零训练的时间,既能加载模型完成手写数字预测,也能继续训练、调整参数或对比不同保存方式,便于理解模型结构与权重关系。对希望降低环境与训练成本、专注模型原理和实验对比的读者较为友好,可快速搭建自己的实验基线。
1. 从一张28×28的灰度图开始:这套MNIST训练资源解决什么问题
拿到一份MNIST数据集训练手写数字识别模型的代码包,最怕的是只看到训练脚本,却不知道前馈神经网络从数据到权重文件怎么串起来。这套资源把数据加载、训练、推理和训练好的模型文件放在同一条链路上,目标很明确:用MNIST数据集跑通手写数字识别,先把前馈神经网络这条基线打牢,再谈换网络。适合刚接触深度学习、需要一份能直接复现的课程设计或小工具的人,也适合熟手拿它当回归测试基线。机器学习项目里最耗时间的往往不是模型结构,而是输入形状、标签格式、权重加载这些细节,这份资源的价值就在把这些细节固定下来。
2. 数据和网络先对齐:前馈神经网络在MNIST上的输入输出怎么定
2.1 为什么先用前馈神经网络,而不是一上来堆卷积
MNIST的每张图是28×28灰度图,展开后是784维向量,类别是0到9共10类。这个规模很适合先验证一条最小闭环:数据能不能读进来、张量形状对不对、损失函数接得对不对、权重能不能保存再加载。前馈神经网络的结构简单,输入层784个节点,中间两三个全连接隐藏层,输出层10个节点,参数规模在二十多万到几十万之间,普通CPU也能在几分钟内跑完一轮训练。对课程设计、小工具原型、教学演示来说,这个速度足够反复试错。
卷积网络在MNIST上通常能拿到更高的准确率,但它引入卷积核大小、步长、填充、池化、通道数等变量,一旦结果不对,排查面会变大。我一般会先用前馈神经网络把数据管线和训练循环跑通,确认测试集准确率能稳定在97%以上,再决定要不要换卷积网络。常见做法是保留同一套数据加载、同一套评估函数,只替换模型定义,这样对比才有意义。前馈网络不是最终答案,但它是一块可靠的底座,后面换模型时,数据预处理和权重保存逻辑不用重写。
这里要区分“模型输入”和“数据形状”。PyTorch读出来的单张MNIST是[1, 28, 28],而全连接层需要[batch, 784]。有人喜欢在数据 transforms 里加transforms.Lambda(lambda x: x.view(-1)),有人喜欢在模型第一层写nn.Flatten()。两种都行,但训练和推理必须一致。我倾向于把展平放在模型里,数据侧保留图像形状,这样推理时传单张图片不容易漏掉维度。
2.2 数据管线:张量形状、归一化和标签格式
先看数据加载。下面的代码把训练集和测试集读进来,转成张量,做MNIST常用的均值和标准差归一化,并打印一个batch的形状和数值范围。batch_size先设128,num_workers设0,是为了在Windows、Notebook和没有多进程保护的环境里少出问题。
import torch from torch.utils.data import DataLoader from torchvision import datasets, transforms # ToTensor 会把 PIL 的 28x28 灰度图转成 [1, 28, 28] 的 float 张量,数值缩到 [0, 1] # Normalize 使用 MNIST 训练集常见统计量:均值 0.1307,标准差 0.3081 transform = transforms.Compose([ transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,)) ]) train_ds = datasets.MNIST( root="./data", train=True, download=True, transform=transform ) test_ds = datasets.MNIST( root="./data", train=False, download=True, transform=transform ) train_loader = DataLoader(train_ds, batch_size=128, shuffle=True, num_workers=0) test_loader = DataLoader(test_ds, batch_size=256, shuffle=False, num_workers=0) x, y = next(iter(train_loader)) print("x shape:", x.shape) # [128, 1, 28, 28] print("x min/max:", x.min().item(), x.max().item()) print("y shape:", y.shape) # [128] print("y first 8:", y[:8])这段代码里,root="./data"是数据目录,第一次运行会下载四个idx文件,之后直接读本地。train=True/False区分60000张训练图和10000张测试图。ToTensor()负责把像素从0到255缩到0到1,并把通道维度提前。Normalize不是必须,但训练和推理要统一:训练用了归一化,推理也必须用同样的均值和标准差,否则输入分布偏移,预测结果会变得很奇怪。标签是0到9的整数,CrossEntropyLoss直接接收这种格式,不需要做one-hot。很多人在这里翻车,是把标签转成one-hot后又塞给交叉熵,结果形状对不上。
batch_size=128是个折中值:太小会让梯度噪声大,太大会让小显存环境吃紧。shuffle=True只对训练集开,测试集保持顺序,方便后面按索引核对样本。num_workers=0不是性能最优,但最稳。如果确定在Linux下跑,可以改成2或4,同时注意DataLoader的随机种子问题。
2.3 网络结构参数:784-256-128-10的每一层在做什么
网络定义如下:第一层全连接把784维映射到256维,ReLU激活,Dropout丢弃一部分神经元;第二层全连接把256维压到128维,再ReLU和Dropout;最后一层输出10维logits。隐藏层宽度和Dropout比例是这份基线里最值得调的两个参数。
import torch.nn as nn class MLP(nn.Module): def __init__(self, hidden1=256, hidden2=128, dropout=0.2): super().__init__() self.net = nn.Sequential( nn.Flatten(), # [B, 1, 28, 28] -> [B, 784] nn.Linear(28 * 28, hidden1), # 784 -> 256 nn.ReLU(), nn.Dropout(dropout), nn.Linear(hidden1, hidden2), # 256 -> 128 nn.ReLU(), nn.Dropout(dropout), nn.Linear(hidden2, 10) # 128 -> 10,输出 logits ) def forward(self, x): return self.net(x) model = MLP(hidden1=256, hidden2=128, dropout=0.2) print(model)nn.Flatten()默认从第1维开始展平,保留batch维,所以输入[128, 1, 28, 28]会变成[128, 784]。第一层全连接参数量是784×256+256,第二层是256×128+128,输出层是128×10+10,加起来约23.5万。这个量级在CPU上训练也不会太慢。Dropout(0.2)表示训练时每个神经元有20%概率被临时丢弃,推理时自动关闭。它的作用是抑制过拟合,但如果数据管线本身有问题,Dropout会把现象掩盖得更隐蔽,所以排查阶段可以先把dropout设0。
输出层不能加Softmax,因为CrossEntropyLoss内部已经做了log-softmax和负对数似然。如果模型里再加一次Softmax,等于做了两次变换,梯度会变弱,训练容易卡住。预测时用argmax(dim=1)取最大logit对应的类别即可;需要置信度时,再对logits做一次softmax。
2.4 损失函数、优化器和初始化:先把默认值跑稳
优化器先用Adam,学习率1e-3,损失函数用交叉熵。权重初始化用PyTorch默认的Kaiming均匀初始化通常就够,不需要一上来手写复杂的初始化。下面的代码把设备、损失函数、优化器准备好,并跑一次前向,确认输出形状是[batch, 10]。
device = torch.device("cuda" if torch.cuda.is_available() else "cpu") model = MLP(hidden1=256, hidden2=128, dropout=0.2).to(device) criterion = nn.CrossEntropyLoss() optimizer = torch.optim.Adam(model.parameters(), lr=1e-3) # 取一个 batch 做一次前向,确认形状和数值范围 x, y = next(iter(train_loader)) x, y = x.to(device), y.to(device) logits = model(x) loss = criterion(logits, y) print("logits shape:", logits.shape) # [128, 10] print("loss:", loss.item())device会自动选GPU或CPU,后面所有张量都要.to(device),模型也要.to(device),两边设备不一致会直接报错。lr=1e-3是Adam在MNIST前馈网络上的常用起点;如果loss出现剧烈震荡,先降到3e-4或1e-4。CrossEntropyLoss接收未归一化的logits和整数标签,不要传概率。前向输出的形状必须是[batch, 10],如果变成[batch, 1, 10]或[10, batch],后面的损失计算和argmax都会出问题。
这一章把数据、网络、损失三件事对齐后,训练脚本才有意义。参数不需要一次调到最优,先把hidden1=256、hidden2=128、dropout=0.2、lr=1e-3这组默认值跑通,再根据验证集表现微调。常见做法是每次只改一个参数,记录测试集准确率和loss曲线,不要同时改网络宽度、学习率和batch size,否则结果变好变坏都说不清原因。
3. 训练脚本落地:从60000张训练图到可复用的权重文件
3.1 训练循环的五个固定动作:前向、损失、反向、更新、清零
训练循环看起来简单,真正容易错的是顺序和累计方式。每个batch固定做五件事:梯度清零、前向计算、计算损失、反向传播、更新参数。顺序不能乱,尤其是optimizer.zero_grad()要放在反向传播之前,否则梯度会累加。下面的代码把单轮训练封装成函数,返回平均loss和准确率。
def train_one_epoch(model, loader, optimizer, criterion, device): model.train() # 开启 Dropout 和 BatchNorm 的训练行为 total_loss = 0.0 correct = 0 total = 0 for x, y in loader: x, y = x.to(device), y.to(device) optimizer.zero_grad() logits = model(x) loss = criterion(logits, y) loss.backward() optimizer.step() # loss.item() 是 batch 平均损失,乘 batch 大小还原成总损失 total_loss += loss.item() * x.size(0) correct += (logits.argmax(dim=1) == y).sum().item() total += x.size(0) return total_loss / total, correct / totalmodel.train()不能省,它决定Dropout是否生效。loss.item()取的是当前batch的平均损失,乘以x.size(0)才是这个batch的总损失,最后除以样本总数得到epoch平均损失。准确率用argmax(dim=1)取预测类别,再和y逐元素比较。这里的correct是整数累加,不要用浮点数累加后再取整,否则最后一位可能对不上。
如果显存不够,优先降batch_size,不要先改网络结构。batch_size=128在大多数环境里够用;如果报CUDA out of memory,降到64或32,同时观察loss是否更抖。学习率是另一个敏感参数:1e-3跑不动时可以试3e-4,不要直接拉到1e-2,否则loss可能先冲高再崩。
3.2 验证集与指标:为什么每轮只看测试集会失真
MNIST官方给了60000张训练图和10000张测试图。严格做法是从训练集里再切一部分做验证集,测试集只在最后评估一次。但课程设计和小工具场景里,常见做法是每轮在测试集上看准确率,用它观察收敛趋势。这样做不是最严谨,但能快速判断模型有没有学起来。下面的评估函数关闭梯度,切换到eval()模式,避免Dropout干扰。
@torch.no_grad() def evaluate(model, loader, criterion, device): model.eval() # 关闭 Dropout,使用推理行为 total_loss = 0.0 correct = 0 total = 0 for x, y in loader: x, y = x.to(device), y.to(device) logits = model(x) loss = criterion(logits, y) total_loss += loss.item() * x.size(0) correct += (logits.argmax(dim=1) == y).sum().item() total += x.size(0) return total_loss / total, correct / total@torch.no_grad()减少显存占用并加快推理,model.eval()关闭Dropout。评估时不要调用loss.backward(),也不要把评估数据用来更新参数。每轮打印训练loss、训练准确率、测试loss、测试准确率,至少能看到两条曲线是否同步。如果训练准确率一直涨、测试准确率不涨,说明过拟合或数据分布有问题;如果两者都停在10%左右,说明模型根本没学到东西,优先查标签、损失函数和输入归一化。
3.3 保存与恢复:state_dict、优化器状态和随机种子
训练完成后要保存权重。建议保存一个字典,里面放model_state、网络超参和输入形状,而不是只保存model.state_dict()。这样推理脚本加载时能重建同样的网络结构,不会因为隐藏层宽度对不上而报错。下面的代码跑10轮训练,每轮评估,最后保存权重文件。
import os # 固定随机种子,方便复现 torch.manual_seed(42) if torch.cuda.is_available(): torch.cuda.manual_seed_all(42) model = MLP(hidden1=256, hidden2=128, dropout=0.2).to(device) criterion = nn.CrossEntropyLoss() optimizer = torch.optim.Adam(model.parameters(), lr=1e-3) for epoch in range(1, 11): train_loss, train_acc = train_one_epoch(model, train_loader, optimizer, criterion, device) test_loss, test_acc = evaluate(model, test_loader, criterion, device) print( f"epoch={epoch:02d} " f"train_loss={train_loss:.4f} train_acc={train_acc:.4f} " f"test_loss={test_loss:.4f} test_acc={test_acc:.4f}" ) os.makedirs("checkpoints", exist_ok=True) torch.save( { "model_state": model.state_dict(), "input_shape": (1, 28, 28), "hidden1": 256, "hidden2": 128, "dropout": 0.2, "normalize_mean": 0.1307, "normalize_std": 0.3081, }, "checkpoints/mnist_mlp.pth", ) print("saved checkpoints/mnist_mlp.pth")保存字典里带了input_shape和归一化参数,是为了让推理脚本不用猜。torch.save保存的是张量字典,不是整个模型对象,加载时更安全,也不依赖训练脚本里的类定义位置。如果继续训练,除了model_state,还要保存optimizer.state_dict()和当前epoch,否则恢复后优化器状态丢失,学习率动量会重新开始。随机种子只固定了PyTorch和CUDA,DataLoader如果开多进程,还需要给worker设置种子;这份基线先用num_workers=0,省掉这层麻烦。
训练日志不要只打印准确率,loss也要看。准确率是离散的,loss能更早反映模型是否在收敛。10轮训练在CPU上通常几分钟到十几分钟,GPU上更快。如果第1轮测试准确率就接近99%,不要高兴太早,先检查测试集有没有被误当成训练集,或者标签有没有泄漏。正常前馈网络在MNIST上从90%多逐步爬到97%到98%左右,是比较常见的曲线。
4. 模型文件直接用:单张推理、批量评估和输入预处理
4.1 加载权重时最容易错的三件事:结构一致、设备一致、模式切换
拿到训练好的模型文件后,第一件事不是直接预测,而是按保存时的超参重建网络,再加载state_dict。下面代码用map_location兼容CPU和GPU,加载后切到eval()。如果加载时报Missing key(s)或Unexpected key(s),先检查hidden1、hidden2是否和保存时一致,再检查是不是把整个模型对象和state_dict混用了。
import torch def load_model(path, device): ckpt = torch.load(path, map_location=device) model = MLP( hidden1=ckpt.get("hidden1", 256), hidden2=ckpt.get("hidden2", 128), dropout=0.0 # 推理阶段不需要 Dropout ) model.load_state_dict(ckpt["model_state"]) model.to(device) model.eval() return model, ckpt device = torch.device("cuda" if torch.cuda.is_available() else "cpu") model, ckpt = load_model("checkpoints/mnist_mlp.pth", device) print("loaded, input_shape:", ckpt.get("input_shape")) print("normalize:", ckpt.get("normalize_mean"), ckpt.get("normalize_std"))map_location=device让权重在CPU上也能加载,不会因为训练时用了GPU而报设备错。dropout=0.0只是重建网络时少建Dropout层,load_state_dict不会受Dropout比例影响,因为Dropout没有参数。model.eval()必须调用,否则推理时Dropout仍然随机丢弃神经元,同一张图多次预测可能得到不同结果。加载成功后先打印超参和归一化参数,后面预处理才有依据。
4.2 单张图片推理:从PIL到归一化张量
单张推理的输入可能来自PIL图片、NumPy数组或截图。下面的函数统一转灰度、缩放到28×28、转张量、加batch维,再做和训练一致的归一化。注意unsqueeze(0)的位置:ToTensor()输出[1, 28, 28],加batch后是[1, 1, 28, 28]。
from PIL import Image from torchvision import transforms import torch.nn.functional as F def predict_image(model, image, device, mean=0.1307, std=0.3081): if image.mode != "L": image = image.convert("L") image = image.resize((28, 28)) x = transforms.ToTensor()(image) # [1, 28, 28],数值 [0, 1] x = x.unsqueeze(0).to(device) # [1, 1, 28, 28] x = transforms.Normalize((mean,), (std,))(x) with torch.no_grad(): logits = model(x) prob = F.softmax(logits, dim=1) conf, pred = prob.max(dim=1) return pred.item(), conf.item(), prob.squeeze(0).cpu() img = Image.open("test_digit.png") pred, conf, prob = predict_image(model, img, device) print("pred:", pred, "confidence:", round(conf, 4)) print("top3:", torch.topk(prob, k=3))image.convert("L")把RGB或RGBA转成单通道灰度。resize((28, 28))要放在ToTensor()之前,因为PIL的resize接收宽高。Normalize的均值和标准差必须和训练时一致;如果训练脚本里没归一化,这里也不能加。MNIST原始图是黑底白字,如果输入是白底黑字,需要先反色,否则模型看到的像素分布完全反过来,预测会偏向某一类。置信度低不代表一定错,但连续多张低置信度,要检查输入预处理和模型是否加载错。
4.3 批量评估与混淆矩阵:别只看一个准确率
单张预测跑通后,用测试集做批量评估。除了总体准确率,还应该看每一类的precision、recall、f1,以及混淆矩阵。MNIST里4和9、3和5、7和1容易混,混淆矩阵能直接看出错在哪。下面的代码收集全部预测结果,用sklearn输出分类报告。
from sklearn.metrics import classification_report, confusion_matrix y_true, y_pred = [], [] model.eval() with torch.no_grad(): for x, y in test_loader: x = x.to(device) logits = model(x) pred = logits.argmax(dim=1).cpu().numpy() y_true.extend(y.numpy().tolist()) y_pred.extend(pred.tolist()) print(classification_report(y_true, y_pred, digits=4)) cm = confusion_matrix(y_true, y_pred) print(cm)classification_report里的support是每类样本数,MNIST测试集每类大约1000张,分布比较均衡。如果某一类recall明显低,先看混淆矩阵里它被错分成什么。比如4被大量分成9,可能是书写风格接近,也可能是归一化或反色没对齐。批量评估的预处理必须和训练完全一致,test_loader里的transform已经包含归一化,所以这里不要再手动加一次。
5. 避坑与排查:MNIST训练链路上的五类翻车现场
5.1 数据与路径类:下载失败、形状错位、标签越界
现象:第一次运行时报连接超时,或者RuntimeError: Dataset not found。原因:download=True需要联网拉取MNIST的四个idx文件,离线环境或缓存目录权限不对就会失败。解决:在能联网的环境先跑一次,把./data/MNIST/raw下的压缩包和解压文件准备好,之后把download=False,确认文件名是train-images-idx3-ubyte、train-labels-idx1-ubyte、t10k-images-idx3-ubyte、t10k-labels-idx1-ubyte。路径不要带中文和空格,Windows下尤其注意反斜杠转义。
现象:模型报mat1 and mat2 shapes cannot be multiplied。原因:输入没有展平,或者展平维度不对。全连接层期望[B, 784],实际传了[B, 1, 28, 28]。解决:在模型第一层加nn.Flatten(),或者在数据侧加transforms.Lambda(lambda x: x.view(-1))。改完后打印一次x.shape,确认是[batch, 784]再进全连接。推理时单张图要unsqueeze(0)补batch维,否则会变成[784],仍然对不上。
现象:标签出现10或负数,交叉熵报Target out of bounds。原因:手动构造标签时索引从1开始,或者把one-hot矩阵直接传给CrossEntropyLoss。解决:MNIST标签就是0到9的整数,CrossEntropyLoss接收[B]的LongTensor。如果手里是one-hot,用argmax(dim=1)转回整数;如果是自己标的数据,检查有没有把背景类标成10。
5.2 训练与收敛类:loss不降、准确率卡住、设备错配
现象:训练几轮后loss一直在2.3左右,准确率停在10%。原因:常见有三种:学习率过大导致震荡,标签和输出错位,或者模型输出加了Softmax后又用交叉熵。解决:先把学习率降到3e-4,把dropout设0,打印一个batch的logits[:2]和y[:2],确认输出是10维、标签是0到9。输出层不要加Softmax,CrossEntropyLoss内部会处理。如果loss还是不降,用一个小数据集过拟合测试:取100张图训练几百步,正常应该能接近100%,过拟合不了说明代码链路有问题。
现象:报Expected all tensors to be on the same device。原因:模型在GPU上,输入张量还在CPU上,或者反过来。解决:统一用device变量,模型.to(device),每个batch的x, y也.to(device)。加载权重时用map_location=device,不要直接torch.load到默认设备。推理脚本里如果用了NumPy数组,转张量后也要.to(device)。
现象:训练准确率很高,测试准确率明显低一截。原因:过拟合,或者训练集和测试集预处理不一致。解决:先检查train_loader和test_loader是不是用了同一个transform;再适当增大Dropout或减小隐藏层宽度。MNIST上前馈网络过拟合不算严重,如果差距很大,优先怀疑测试集被污染或归一化参数不一致。不要用测试集反复调参调到刚好,留一部分训练数据做验证更稳。
5.3 保存与复现类:权重加载失败、结果对不上、随机性失控
现象:load_state_dict报Missing key(s)或size mismatch。原因:保存时用了hidden1=256,加载时建了hidden1=128;或者保存的是整个模型对象,加载时又按state_dict读。解决:保存字典里带上hidden1、hidden2,加载时按字典重建网络。统一用model.state_dict()保存参数,不要torch.save(model)保存整个对象。加载时先打印ckpt.keys(),确认键名是model_state还是state_dict。
现象:同一张图两次预测结果不一样,或者重新训练后准确率差很多。原因:推理时忘了model.eval(),Dropout还在随机丢弃;训练时没有固定随机种子;DataLoader多进程的worker种子没设。解决:推理前固定调用model.eval(),并用torch.no_grad()包住。训练前设torch.manual_seed、torch.cuda.manual_seed_all,需要更严格时把cudnn.deterministic设为True。num_workers>0时给DataLoader传generator和worker_init_fn,或者先保持num_workers=0把结果跑稳。
现象:训练时准确率正常,推理脚本预测全是同一类。原因:预处理没对齐,最常见的是训练用了Normalize,推理没用;或者输入图片是白底黑字,没有反色;或者图片没缩放到28×28。解决:把训练时的mean、std存进权重字典,推理时读出来用。打印推理输入的最小值、最大值、均值和训练集对比,数值范围差太多就说明预处理错了。白底黑字先转灰度再反色,确保和MNIST的黑底白字一致。
6. 把基线做成可验证的小工具:导出ONNX与固定随机种子复现
前馈网络跑通后,我习惯做两件事:固定随机种子,把一次训练结果复现出来;导出ONNX,用另一个运行时验证输出是否一致。固定种子不是玄学,是为了让“这次准确率98.1%,下次97.9%”这种波动有参照。下面的函数把Python、NumPy、PyTorch和CUDA的种子一起固定,训练前调用一次。
import random import numpy as np import torch def seed_everything(seed=42): random.seed(seed) np.random.seed(seed) torch.manual_seed(seed) torch.cuda.manual_seed_all(seed) torch.backends.cudnn.deterministic = True torch.backends.cudnn.benchmark = False seed_everything(42)导出ONNX时给一个固定输入,检查ONNX Runtime的输出和PyTorch输出是否接近。下面的代码用opset_version=12,兼容性比较稳;如果环境里的ONNX Runtime版本较新,也可以试更高的opset,但不要一上来就拉满。
import numpy as np import torch import onnxruntime as ort dummy = torch.randn(1, 1, 28, 28).to(device) torch.onnx.export( model, dummy, "mnist_mlp.onnx", input_names=["input"], output_names=["logits"], opset_version=12, ) sess = ort.InferenceSession("mnist_mlp.onnx", providers=["CPUExecutionProvider"]) onnx_logits = sess.run(None, {"input": dummy.cpu().numpy()})[0] torch_logits = model(dummy).detach().cpu().numpy() print("max abs diff:", np.max(np.abs(onnx_logits - torch_logits)))max abs diff在1e-5到1e-4量级通常可以接受,如果差到1e-2以上,先检查模型是不是还在训练模式,或者导出时输入形状和推理时不一致。ONNX导出后,推理脚本不再依赖PyTorch,适合把模型塞进小工具里。固定种子加ONNX校验,相当于给这条MNIST基线上了双保险:一次训练能复现,一份权重能跨运行时验证。从那以后我每次拿到新的MNIST训练包,都先固定种子跑一遍单张推理,再决定要不要改网络。希望帮到你。
本文还有配套的精品资源,点击获取