PyTorch实战MNIST手写数字识别:从数据加载到模型训练 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, trainTrue, transformtransform, downloadTrue) test_dataset datasets.MNIST(root./data, trainFalse, transformtransform, downloadTrue)这里有两个关键点需要理解。第一trainTrue拿到的是训练集trainFalse拿到的是测试集两者不要混。第二downloadTrue表示如果本地没有数据就自动下载。下载后的文件会放在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_size256, shuffleTrue) test_loader DataLoader(test_dataset, batch_size256, shuffleFalse)batch_size256表示每个批次包含 256 张图片。shuffleTrue表示每个 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(), lr0.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(dim1) 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(dim1) total_correct (preds labels).sum().item() total_count images.size(0) return total_loss / total_count, total_correct / total_count评估指标我选择的是准确率这个指标在 MNIST 上最直观也最容易理解。preds logits.argmax(dim1)表示取出每个样本得分最高的那个类别的索引然后和真实标签对比统计相等的个数。对于 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, trainTrue, transformtransform, downloadTrue) test_dataset datasets.MNIST(root./data, trainFalse, transformtransform, downloadTrue) # 构建数据加载器 train_loader DataLoader(train_dataset, batch_size256, shuffleTrue) test_loader DataLoader(test_dataset, batch_size256, shuffleFalse) # 定义模型 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(), lr0.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(dim1) 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(dim1) 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(fEpoch {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 数据下载失败或下载速度很慢这是初学者最容易遇到的第一道坎。downloadTrue会尝试从数据集官网下载文件如果你的网络不稳定下载过程可能中断导致脚本直接报错。第一次看到网络相关报错的时候先别慌这不是你的代码有问题。解决办法有几个。最简单的就是重新运行脚本PyTorch 会从断点开始重新下载已下载的部分。如果反复失败你可以手动把数据集文件下载到./data目录下并确保文件名和 PyTorch 预期的一致然后再设置downloadFalse。还可以换一个网络环境试试比如单位的网络、公共网络很多时候换个网络就好了。这里也要提醒一个点如果下载好的文件被反复校验失败可能是压缩包不完整。删除./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_size3, padding1), nn.ReLU(), nn.MaxPool2d(2), nn.Conv2d(32, 64, kernel_size3, padding1), 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 是这个习惯最好的练习场养成它之后你后面接触更大规模的数据集会轻松很多。