☰
PyTorch实战MNIST手写数字识别:从数据加载到模型训练
2026/10/17 18:07:40 网站建设 项目流程

MNIST 是入门深度学习绕不开的第一个数据集。一个 28x28 的灰度手写数字图片,从原始像素到模型输出一个 0-9 的类别,这中间经历的每一个环节,几乎就是深度学习项目的最小闭环。很多新手看过理论之后卡在“代码怎么组织”“数据怎么喂给模型”“训练完怎么评估”,这篇文章就用完整的 PyTorch 代码把这个流程走一遍,从数据加载、预处理、模型设计到训练与评估,每一步都会解释为什么这么做。

这篇实战笔记适合几类人:刚看完神经网络理论但不知道代码怎么写的人;想用一个小项目快速验证自己对数据流理解的人;以及准备换框架、拿 MNIST 当基线来跑通流程的人。MNIST 的数据量很小,纯 CPU 也能在几分钟内完成训练,所以环境门槛很低,你不需要一台多好的电脑,先把整个项目闭环跑起来,比什么都重要。

1. 项目概述与整体思路

1.1 MNIST 是什么,为什么要用它

MNIST 是一个手写数字图片数据集,包含 60000 张训练图片和 10000 张测试图片。每张图片都是 28x28 的灰度图,像素值范围是 0 到 255,对应的标签是 0 到 9 之间的一个整数。它解决的问题很简单:给模型看一张图片,让它判断这张图片写的是哪个数字。

这个数据集被用了几十年,现在已经非常“干净”:图片尺寸固定、类别均衡、没有太多噪声。正因为如此,它特别适合用来学习深度学习的基本流程。你不用花大量时间做图像清理和数据增强,可以把精力集中在“数据怎么进模型”“梯度怎么更新”“准确率怎么评估”这些核心环节上。

还有一个容易被忽略的好处:MNIST 的训练速度足够快。一个简单的全连接网络在 CPU 上跑 10 个 epoch 通常只要几分钟,这让你可以反复修改代码、反复实验,而不需要等半天才能看到结果。我见过不少初学者一上来就用很复杂的模型跑大规模数据,结果训练一次要几小时,最后连“为什么模型不收敛”都没耐心排查。先用 MNIST 这样的数据集把流程跑通,才是性价比最高的路径。

1.2 整体数据流与方案选型

这个项目的核心是“数据流”。完整流程可以概括为:原始图片 -> 数据集对象 -> 张量预处理 -> 分批加载器 -> 模型前向传播 -> 损失计算 -> 反向传播 -> 参数更新 -> 测试集评估。很多人把这套流程拆成“数据处理”和“模型训练”两个独立部分,但实际上它们是一个整体。

我在方案选型上选择了 PyTorch,原因很直接:它的张量操作和自动求导机制让你能清楚看到每个环节的数据形态,对理解深度学习底层逻辑很有帮助。你拿到一个 batch 的图片,形状是(256, 1, 28, 28),经过nn.Flatten之后变成(256, 784),再经过线性层变成(256, 10)。每一步的形状变化都清晰可见,这对新手来说非常重要。

如果你更习惯其他框架,这个思路照样适用。核心不在于用哪个库,而在于你能否回答这几个问题:数据加载后是什么类型?模型输入是什么形状?损失函数接受的标签是什么格式?这几个问题一旦搞清楚,换框架只是语法层面的差异。

1.3 环境与依赖准备

这个项目的依赖非常少,核心只有三个:PyTorch、torchvision 和 Python。torchvision 用来下载和加载 MNIST 数据集,PyTorch 提供模型训练所需的自动求导和优化器。

安装命令一般是:

pip install torch torchvision

如果你用的是在线 Notebook 环境,通常已经预装了这些库,可以直接开始。本地环境的话,CPU 版本就能满足需求,因为 MNIST 图片实在太小了,10 层以内的简单网络在 CPU 上跑都毫无压力。GPU 可以在后续扩展到更大数据集时再考虑,现阶段不要让环境配置成为你开始学习的阻碍。

我把主要依赖整理成了一张表:

依赖项作用版本建议
Python脚本运行环境3.8 及以上
PyTorch张量计算、自动求导、模型训练2.x 均可
torchvision数据集加载与图像预处理工具随 PyTorch 版本配套安装

2. 数据流详解与预处理

2.1 数据集的加载与目录结构

MNIST 数据集的加载在 PyTorch 中非常简单,核心逻辑就是调用torchvision.datasets.MNIST。它会在第一次运行的时候把数据下载到本地,之后再次运行会直接读取已经下载好的文件,不会重复下载。

from torchvision import datasets, transforms transform = transforms.Compose([ transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,)) ]) train_dataset = datasets.MNIST(root="./data", train=True, transform=transform, download=True) test_dataset = datasets.MNIST(root="./data", train=False, transform=transform, download=True)

这里有两个关键点需要理解。第一,train=True拿到的是训练集,train=False拿到的是测试集,两者不要混。第二,download=True表示如果本地没有数据就自动下载。下载后的文件会放在root指定的./data目录下,里面包含train-images-idx3-ubyte.gz和train-labels-idx1-ubyte.gz这类文件。

我建议你把root目录固定下来,比如./data,这样多次运行脚本时不会每次都下载。实际项目中,数据文件的组织方式也会直接影响代码的可维护性,尽早养成“数据固定放在某个目录”的习惯会省掉很多麻烦。

2.2 图像转换与标准化

原始 MNIST 图片是 PIL Image 对象,像素值范围是 0 到 255。直接把这种数据丢给神经网络不是不行,但效果通常不够好。transforms.ToTensor()会做两件事:把图片转成 PyTorch 的 Tensor,同时把像素值缩放到 0 到 1 之间。转换之后,一张 28x28 的灰度图会变成一个形状为(1, 28, 28)的张量,多出来的第一维是通道数。

transforms.Normalize((0.1307,), (0.3081,))这一步是标准化。它会对每个像素做(x - mean) / std的运算,让数据分布接近均值为 0、标准差为 1 的标准正态分布。0.1307和0.3081是 MNIST 数据集全体像素的均值和标准差,这不是随便写的,而是提前统计出来的固定值。

为什么要标准化?我常用一个类比:如果数据里的特征量级差异很大,模型优化时就像在高低不平的地面上走路,每一步都需要适应不同坡度;标准化之后,地面被大致“铺平”了,梯度下降的路径会顺畅很多。对于 MNIST 这种灰度图,不做标准化也能训练,但做了之后收敛速度更快,训练过程也更稳定。

2.3 DataLoader 批次划分与乱序

数据集准备好之后,下一步是用DataLoader把它包装成一个可以分批读取的加载器。这是深度学习中“数据流”的关键一环。你不可能一次性把 60000 张图片全部塞进模型计算,一方面内存可能不够,另一方面大 batch 的梯度更新也不一定稳定。所以需要把数据切成一个个小批次。

from torch.utils.data import DataLoader train_loader = DataLoader(train_dataset, batch_size=256, shuffle=True) test_loader = DataLoader(test_dataset, batch_size=256, shuffle=False)

batch_size=256表示每个批次包含 256 张图片。shuffle=True表示每个 epoch 开始前把数据顺序打乱,这可以防止模型学到数据集中固定的排列顺序,从而提高泛化能力。测试集不需要打乱,因为评估结果不应该受数据顺序影响。

我见过不少初学者好奇“为什么训练的时候要 shuffle”。简单来说,如果数据天然有一定顺序,比如前 1000 张都是数字 0,模型在连续几轮都只看到数字 0,梯度就会偏得很厉害,甚至出现“刚学完 0 就忘了 1”的灾难性遗忘。打乱之后,每个 batch 里都混合了不同类别,梯度更新方向才更接近全局方向。

3. 模型设计

3.1 为什么先用全连接网络

MNIST 的经典做法可以有很多种:全连接网络、卷积神经网络、甚至现在的各种轻量模型。但对于入门实战,我强烈建议先用全连接网络,也就是多层感知机。原因只有一个:它能让你把注意力集中在训练流程本身,而不是模型结构带来的额外复杂度。

卷积网络确实在图像任务上表现更好,但它引入了卷积核、通道数、池化、特征图等一堆新概念。如果这些概念还没吃透,同时又遇到梯度不收敛的问题,你会很难判断到底是训练代码写错了,还是模型结构设计有问题。全连接网络没有这些负担,输入是 784 维向量,中间经过若干线性变换和激活函数,输出是 10 维向量,每一步都可以手算出来。

等你在 MNIST 上把全连接网络跑通,再切换成卷积网络做对比,你会看到识别率明显提升,这时候你对“模型表达能力”这个概念会有非常直观的体会。反过来如果你一上来就写卷积网络,大概率只是把别人的代码抄了一遍,内部发生了什么还是一头雾水。

3.2 模型结构定义与参数计算

我用的模型只有三个线性层,中间夹着 ReLU 激活函数。先把图片拉直成 784 维向量,然后依次通过 512 个神经元、256 个神经元,最后输出 10 个数值。这 10 个数值就是模型对每个数字类别的打分,分数最高的那个类别就是预测结果。

import torch.nn as nn class MLP(nn.Module): def __init__(self): super().__init__() self.flatten = nn.Flatten() self.net = nn.Sequential( nn.Linear(28 * 28, 512), nn.ReLU(), nn.Linear(512, 256), nn.ReLU(), nn.Linear(256, 10) ) def forward(self, x): x = self.flatten(x) return self.net(x)

很多教程会直接给模型结构,不告诉你里面到底有多少参数。我建议你自己动手算一遍,这对理解线性层非常有帮助。

层权重形状偏置形状参数量
第一层(512, 784)(512,)401920
第二层(256, 512)(256,)131328
第三层(10, 256)(10,)2570
合计--535818

计算方法是:权重参数数量等于输入维度乘以输出维度,再加上偏置数量。第一层就是 784 乘以 512 再加 512,结果是 401920。整个模型总共有约 53 万个参数,对一个只有 60000 张图片的数据集来说,这个容量其实已经不算小了。这也解释了为什么后面可能要加 Dropout 或者正则化,因为模型容量足够大时,很容易把训练集背下来。

3.3 损失函数与优化器选择

分类任务最常用的损失函数是交叉熵损失。PyTorch 里的nn.CrossEntropyLoss是一个组合函数,它把 Softmax 和负对数似然损失合在了一起。这意味着模型的最后一层不需要自己加 Softmax,直接输出原始分数 logits 就行。

criterion = nn.CrossEntropyLoss()

这里有个非常容易踩的坑:CrossEntropyLoss接收的标签是整数形式的类别编号,不是 one-hot 编码。也就是说,对于一张手写数字 3 的图片,标签就是一个整数3,而不是[0, 0, 0, 1, 0, 0, 0, 0, 0, 0]。PyTorch 会自动帮你把整数标签转换成内部需要的 one-hot 形式。如果你在别的地方见过手动转 one-hot 的代码,在这里不需要。

优化器我选择 Adam,学习率设置为 0.001。Adam 的优势在于它对学习率不那么敏感,收敛速度也比较快,非常适合新手。如果你用 SGD(随机梯度下降),可能需要调学习率和动量项,这些细节会在你还没搞清楚模型为什么没收敛的时候增加额外的排查成本。

import torch.optim as optim optimizer = optim.Adam(model.parameters(), lr=0.001)

等你想深入理解优化器时,可以把 Adam 换成 SGD 加动量对比一下效果,这是一个很好的进阶实验。但在入门阶段,我建议先把训练流程跑通,再谈调参。

4. 训练与评估完整实现

4.1 训练主循环核心步骤

训练过程本质上是一个循环:取一个批次的数据,计算模型输出,计算损失,反向传播,更新参数,再取下一个批次。这个循环外面再套一层 epoch 循环,表示“把整个训练集从头到尾看一遍”。

def train_one_epoch(model, loader, criterion, optimizer): model.train() total_loss = 0.0 total_correct = 0 total_count = 0 for images, labels in loader: optimizer.zero_grad() logits = model(images) loss = criterion(logits, labels) loss.backward() optimizer.step() total_loss += loss.item() * images.size(0) preds = logits.argmax(dim=1) total_correct += (preds == labels).sum().item() total_count += images.size(0) return total_loss / total_count, total_correct / total_count

有几个细节需要特别说明。第一,model.train()这行不是摆设,它会把模型切换到训练模式。如果你模型里没有 Dropout 或 BatchNorm,这个调用看起来没什么用,但我会一直保留,因为这是一个好习惯。第二,optimizer.zero_grad()必须在每次反向传播之前清空梯度,否则上一次的梯度会累加到这一次上,导致参数更新方向错误。

关于损失累加,我用的是loss.item() * images.size(0)而不是直接累加loss.item()。原因是最后一个批次的样本数可能不足 256,如果直接累加 loss 再除以总批次数,最后的平均损失会偏高。乘以批次大小再除以总样本数,才是真正的加权平均。这个细节虽然对最终准确率影响不大,但能让你打印出来的 loss 曲线更准确。

4.2 验证与测试评估细节

测试评估和训练在代码上有几个重要区别。首先,评估时不需要计算梯度,因为你不更新参数。PyTorch 里可以用torch.no_grad()包裹评估代码,这能省下大量内存和计算时间。其次,模型要切换到评估模式,也就是model.eval(),这会让 Dropout 和 BatchNorm 的表现与训练时不同。

def evaluate(model, loader, criterion): model.eval() total_loss = 0.0 total_correct = 0 total_count = 0 with torch.no_grad(): for images, labels in loader: logits = model(images) loss = criterion(logits, labels) total_loss += loss.item() * images.size(0) preds = logits.argmax(dim=1) total_correct += (preds == labels).sum().item() total_count += images.size(0) return total_loss / total_count, total_correct / total_count

评估指标我选择的是准确率,这个指标在 MNIST 上最直观,也最容易理解。preds = logits.argmax(dim=1)表示取出每个样本得分最高的那个类别的索引,然后和真实标签对比,统计相等的个数。对于 MNIST 这样的均衡分类数据集,准确率基本够用。但在实际项目中,如果类别不平衡,你还要考虑精确率、召回率、F1 分数等指标。

我习惯在每个 epoch 结束后都跑一次测试集评估,而不是只在训练完全结束后跑一次。这样你能看到训练和测试准确率之间的差距如何变化,这是判断过拟合的重要依据。

4.3 完整代码合辑

把上面这些片段整合到一起,就是一个可以直接运行的完整脚本。我建议你先复制这段代码跑通,然后在此基础上做修改。

import torch import torch.nn as nn import torch.optim as optim from torch.utils.data import DataLoader from torchvision import datasets, transforms # 数据预处理 transform = transforms.Compose([ transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,)) ]) # 加载数据集 train_dataset = datasets.MNIST(root="./data", train=True, transform=transform, download=True) test_dataset = datasets.MNIST(root="./data", train=False, transform=transform, download=True) # 构建数据加载器 train_loader = DataLoader(train_dataset, batch_size=256, shuffle=True) test_loader = DataLoader(test_dataset, batch_size=256, shuffle=False) # 定义模型 class MLP(nn.Module): def __init__(self): super().__init__() self.flatten = nn.Flatten() self.net = nn.Sequential( nn.Linear(28 * 28, 512), nn.ReLU(), nn.Linear(512, 256), nn.ReLU(), nn.Linear(256, 10) ) def forward(self, x): x = self.flatten(x) return self.net(x) model = MLP() criterion = nn.CrossEntropyLoss() optimizer = optim.Adam(model.parameters(), lr=0.001) # 训练一个 epoch def train_one_epoch(model, loader, criterion, optimizer): model.train() total_loss = 0.0 total_correct = 0 total_count = 0 for images, labels in loader: optimizer.zero_grad() logits = model(images) loss = criterion(logits, labels) loss.backward() optimizer.step() total_loss += loss.item() * images.size(0) preds = logits.argmax(dim=1) total_correct += (preds == labels).sum().item() total_count += images.size(0) return total_loss / total_count, total_correct / total_count # 评估 def evaluate(model, loader, criterion): model.eval() total_loss = 0.0 total_correct = 0 total_count = 0 with torch.no_grad(): for images, labels in loader: logits = model(images) loss = criterion(logits, labels) total_loss += loss.item() * images.size(0) preds = logits.argmax(dim=1) total_correct += (preds == labels).sum().item() total_count += images.size(0) return total_loss / total_count, total_correct / total_count # 训练主循环 epochs = 10 for epoch in range(1, epochs + 1): train_loss, train_acc = train_one_epoch(model, train_loader, criterion, optimizer) test_loss, test_acc = evaluate(model, test_loader, criterion) print(f"Epoch {epoch:02d} | train_loss {train_loss:.4f} | train_acc {train_acc:.4f} | test_loss {test_loss:.4f} | test_acc {test_acc:.4f}") # 保存模型参数 torch.save(model.state_dict(), "mnist_mlp.pt")

在这个默认配置下,10 个 epoch 之后测试准确率通常能到 97% 左右。如果你发现自己的准确率明显偏低,别急着改模型,先检查数据预处理、损失函数和优化器是不是和代码一致。

4.4 流程图:从数据流到模型评估

标题里说了含流程图,这里我用文字形式把这个流程画出来。它对应了上面代码的执行顺序,也对应了任何深度学习训练任务的最小骨架。

MNIST原始图片 | v 加载数据集并检查样本结构 | v ToTensor:像素值缩放到 0-1 | v Normalize:标准化到标准正态分布 | v DataLoader:分批加载 + 随机打乱 | v 定义模型:784 -> 512 -> 256 -> 10 | v 前向传播:得到预测分数 logits | v 计算损失 CrossEntropyLoss | v 反向传播:计算梯度 | v 优化器更新:Adam 更新参数 | v 循环到完成所有 epoch | v 测试集评估:准确率 / 损失 | v 保存模型参数

这张图看起来很简单,但它揭示了一个重要事实:大多数训练代码出问题,都出在“数据形态和模型期望不一致”或者“训练和评估模式没有正确切换”这两类问题上。你以后做任何深度学习项目,都可以用这张图对照排查。

5. 常见问题与排查技巧

5.1 数据下载失败或下载速度很慢

这是初学者最容易遇到的第一道坎。download=True会尝试从数据集官网下载文件,如果你的网络不稳定,下载过程可能中断,导致脚本直接报错。第一次看到网络相关报错的时候,先别慌,这不是你的代码有问题。

解决办法有几个。最简单的就是重新运行脚本,PyTorch 会从断点开始重新下载已下载的部分。如果反复失败,你可以手动把数据集文件下载到./data目录下,并确保文件名和 PyTorch 预期的一致,然后再设置download=False。还可以换一个网络环境试试,比如单位的网络、公共网络,很多时候换个网络就好了。

这里也要提醒一个点:如果下载好的文件被反复校验失败,可能是压缩包不完整。删除./data目录重新下载,往往比继续排查更省时间。

5.2 训练集准确率很高,测试集准确率明显偏低

这是典型的过拟合信号。训练准确率能到 99%,测试准确率只有 97%,这个差距在 MNIST 上很常见,因为模型容量比较大。你的模型可能记住了训练集里某些无关的噪声特征,而不是真正学到了手写数字的通用规律。

解决办法可以从几个方向入手。第一,增加训练数据多样性,MNIST 本身是静态数据集,你可以做小幅旋转、平移、加噪声等数据增强。第二,在模型中加 Dropout 层,比如在第二个线性层后面加一个nn.Dropout(0.3)。第三,降低模型容量,比如把 512 个神经元改成 256 个。第四,加入 L2 正则化,也就是在优化器中设置weight_decay。

对于入门阶段,我建议你先体验一下“加 Dropout 之后测试准确率是否变化”这个实验。它比单纯背结论更有价值。

5.3 显存不足或者训练速度很慢

MNIST 图片尺寸很小,正常情况下不会出现显存不足的问题。如果你是用 GPU 训练,最常见的原因是 batch size 设置得太大,比如一次性把 60000 张图片全部塞进去。但 MNIST 的 60000 张图片其实也不大,真遇到显存不足,一般是模型本身定义得太复杂,或者数据处理过程有冗余。

如果你在 CPU 上训练,感觉速度很慢,可以先确认是否每个 epoch 都要跑几分钟。MNIST 搭配这个 MLP 模型,10 个 epoch 通常在几分钟内能完成。如果特别慢,可以检查一下是不是装了 GPU 版 PyTorch 但没有真正使用 GPU,这种情况代码会在 CPU 上运行,速度并不快。

想要用 GPU 加速,需要把模型和每个 batch 的数据都移到 GPU 上。代码大致是这样:

device = torch.device("cuda" if torch.cuda.is_available() else "cpu") model.to(device) for images, labels in loader: images, labels = images.to(device), labels.to(device)

注意输出张量也要在同一个设备上才能计算准确率,否则会报设备不匹配的错误。

5.4 训练损失不下降或者准确率一直很低

这种情况首先要区分是“模型没学会”还是“代码有 bug”。如果损失值几乎不变化,或者准确率一直停留在 10% 左右,大概率是下面这几个问题之一。

第一个常见问题是学习率设置不合理。学习率太大,损失会在一个范围内震荡不收敛;学习率太小,收敛速度极慢,看起来就像没在训练。你可以先打印出前几个 batch 的损失值,看是否在缓慢下降。第二个问题是数据和模型不匹配。比如模型输入期望的是(batch, 784),但你传给它的数据是(batch, 1, 28, 28),这时候全连接层会报错。如果你没看到报错,但准确率很低,可以检查一下是不是忘了flatten。

我还遇到过一种情况:训练集和测试集使用了不同的预处理方式,比如训练时用了标准化,测试时忘了用,导致测试数据分布和训练数据不一致,准确率自然会崩。排查这种问题最好的办法是写一小段代码,单独打印一个 batch 的数据范围和标签,确保训练和评估走的是同一条数据通路。

6. 后续扩展建议

6.1 从 MLP 到 CNN 的小步改进

跑通 MLP 之后,下一步建议试试卷积网络。你不需要马上写一个很深的网络,先用一个简单的结构:一个卷积层加一个池化层,再接一个全连接层。把模型的参数量和准确率对比一下,你会更深刻地体会到“为什么图像任务更适合卷积结构”。

一个可以用于对比的简单结构如下:

class SimpleCNN(nn.Module): def __init__(self): super().__init__() self.conv = nn.Sequential( nn.Conv2d(1, 32, kernel_size=3, padding=1), nn.ReLU(), nn.MaxPool2d(2), nn.Conv2d(32, 64, kernel_size=3, padding=1), nn.ReLU(), nn.MaxPool2d(2) ) self.fc = nn.Linear(64 * 7 * 7, 10) def forward(self, x): x = self.conv(x) x = x.flatten(1) return self.fc(x)

这个模型在 MNIST 上很快就能达到 99% 左右的准确率。你可能会惊讶于“模型更复杂反而参数更少且效果更好”,这正是卷积层权值共享带来的优势。

6.2 把 MNIST 练好的关键习惯

我在反复跑 MNIST 的过程里,最深的体会是:不要小看这个“入门级”数据集。它虽然简单,但完整覆盖了数据加载、预处理、模型定义、训练循环、评估、保存模型这整个链路。以后你做任何真实项目,底层逻辑都差不多,区别只是数据更脏、模型更复杂、训练时间更长。

所以我建议你花点时间把这个项目“吃透”,而不是“跑通就完事”。你可以试着做几个小实验:把 batch size 从 256 改成 16 和 1024,观察训练速度和最终准确率的变化;把学习率从 0.001 改成 0.1,观察损失是否发散;把测试集评估去掉,看看训练过程会有什么不一样。这些实验加起来可能只需要半小时,但对理解深度学习训练机制的价值非常大。

最后再分享一个小技巧:每次跑实验前,先打印一个 batch 的数据形状和标签形状。这个习惯看着多余,但它能帮你排除掉大量“数据维度不匹配”的问题。我在实际项目中排查过无数个训练异常,最后发现有一半以上都是从“数据形状根本不对”开始的。MNIST 是这个习惯最好的练习场,养成它之后,你后面接触更大规模的数据集会轻松很多。

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

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

立即咨询