PyTorch实现简单分类网络:MNIST手写数字识别全流程教程
写个简单的分类网络没想到成了我这些年被问得最多的一篇教程。倒不是说它有多高深而是很多刚接触PyTorch的朋友恰恰需要这样一个“足够简单又不敷衍”的完整例子——把全流程跑通再回头啃文档效率高得多。这篇东西的核心思路就一条用最快的速度建立一个可运行的分类模型同时把模型结构、训练循环、数据流水线和踩坑点都交代清楚让看完的人能脱离示例自己动手改。1. 整体设计思路拆解动手写代码前得先把几个关键问题想明白选型都是有理由的不是为了炫技。1.1 为什么用全连接网络而不是CNN一个常见误解是“图像分类不都得用卷积吗”。实际上数据集是MNIST28x28灰度手写数字用全连接网络完全跑得动参数量大概在20万左右CPU上几分钟就能训完一轮。它的好处第一是简单没有卷积核、池化、特征图这些概念需要理解第二是方便调参网络深了不好调试。当你理解了全连接网络怎么反向传播再去换成卷积网络只是替换网络定义的问题。做图像分类直接上CNN也没有错但作为教学案例全连接网络把“深度学习是怎么回事”讲得更直白。1.2 二分类问题如何扩展成多分类场景MNIST是10个类别本质上是多分类。但很多初学者拿到的第一个需求是“判断是不是某类东西”比如猫狗二分类。我建议直接用多分类写法类别数设成2原因是多分类代码可以无缝扩展十类能做两类也能做代码结构不用动。真正要理解的分界点是二分类可以用 Sigmoid BCELoss输出一个概率值多分类用 Softmax CrossEntropyLoss输出每个类别的概率分布多标签分类一张图有多个类别是另一个领域用 Sigmoid BCEWithLogitsLoss这里选了多分类的标准组合实用性最强。1.3 固定随机种子的重要性第一次跑通代码后你会发现每次运行结果略有不同。这是因为参数初始化是随机的。实验复现的关键操作是设置随机种子import torch import numpy as np import random def set_seed(seed42): torch.manual_seed(seed) torch.cuda.manual_seed_all(seed) np.random.seed(seed) random.seed(seed)这会保证每次训练初始条件一致。很多研究者还会固定PyTorch本身的确定性行为但入门阶段设好上面这段就够了。数据加载时要不要设置shuffle、初始化方式怎么选都可以在后面做控制变量实验时调节。1.4 训练/验证/测试三份数据的划分逻辑常见做法是拿到数据集后分成训练集和测试集但这是不严谨的。用简单网络跑分类划分三份数据很有必要训练集用于更新模型权重验证集用于调整超参数学习率、层数、神经元数量和早停判断测试集用于最终评估模型这个集合绝不能在训练时被偷看MNIST的默认划分是60000张训练 10000张测试。实操中常用torch.utils.data.random_split从训练集切出5000张做验证。虽然也可以直接拿官方划分的测试集当验证集但最好养成分三份的习惯后续换数据集不踩坑。2. 核心组件逐个拆解真正动手写代码前需要把几个核心组件理解透每个都有不可替代的作用。2.1 Dataset 与 DataLoader 的分工机制Dataset负责“取数据”按索引返回样本和标签DataLoader负责“组织数据”分批、打乱、并行加载。这两个概念分离带来两个好处内存能承受大数据集不用一次性全部载入且可以灵活组合数据增强、采样器等操作。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)第一行ToTensor()把PIL图像变成[0,1]的Tensor第二行Normalize用MNIST数据集的均值和标准差做标准化。很多人会问这两个数字怎么来的是官方统计出来的全量均值。2.2 为什么要用 DataLoader 分批训练深度学习训练在工程实现上几乎从不一次性把所有数据塞进显卡。假设MNIST单张图是28x28784个像素6万张一次载入用全连接网络计算显存直接吃满。分批训练Mini-batch让显卡可以分批吃数据同时更新频率增加收敛也更稳定。from torch.utils.data import DataLoader train_loader DataLoader(train_dataset, batch_size64, shuffleTrue)batch_size 的选择影响训练稳定性和显存占用。64在公开数据集上是个安全默认值。不用去读几千篇文章自己跑一组实验对比一下感受会更直观。2.3 模型继承 nn.Module 的背后逻辑定义一个网络要继承nn.Module这不是“格式要求”而是一整套自动求导、参数注册、训练/评估模式切换机制的基础。nn.Module会递归登记所有子模块和参数于是model.parameters()能统一返回全网络参数model.train()和model.eval()能控制BatchNorm、Dropout等层的工作状态。典型网络定义import torch.nn as nn class SimpleNN(nn.Module): def __init__(self, input_size784, hidden_size128, num_classes10): super(SimpleNN, self).__init__() self.fc1 nn.Linear(input_size, hidden_size) self.relu nn.ReLU() self.fc2 nn.Linear(hidden_size, num_classes) def forward(self, x): x x.view(x.size(0), -1) x self.fc1(x) x self.relu(x) x self.fc2(x) return x__init__里定义层forward里定义数据流向。损失函数部分很关键CrossEntropyLoss内部已经包含了 Softmax 操作所以网络最后一层只需要输出原始分数Logits不需要额外加 Softmax。2.4 激活函数为什么选的 ReLU早期神经网络常用 Sigmoid 或 Tanh但反向传播时会出现梯度消失——网络加深后误差信号从输出层传回前面层时被层层缩放参数几乎不再更新。ReLU 在正区间梯度恒为1能有效缓解这个问题而且计算便宜。ReLU也有缺点比如负区间输出恒为0部分神经元可能“死掉”。入门阶段不必纠结这个先把流程跑通。后面换 LeakyReLU 或 GELU 可以当作一次调参实验。3. 数据与模型实现细节基础理论到位下面是实践中会遇到的具体实现问题。3.1 数据预处理标准化和增强的取舍图像标准化通常是每通道减去均值除以标准差。对MNIST来说灰度图只有一个通道。某些场景要加数据增强操作例如随机裁剪、旋转、水平翻转等能提升泛化能力。但做增强前先要清醒认识到MNIST本身是规范数据不需要太多增强换成真实场景数据后增强的效果才会得到体现。标注downloadTrue时要注意网络状况。如果下载慢或失败建议手动从官网下载四个压缩包放进./data/MNIST/raw/目录。版本不匹配时torchvision会报RuntimeError找不到文件或md5校验失败遇到这个情况优先检查压缩包放置路径和文件完整性。3.2 损失函数的选择与数学含义分类问题标准的损失函数是交叉熵Cross Entropy。从信息论角度理解它衡量“用预测分布去编码真实标签所需的信息量”预测越接近真实损失越低。PyTorch 的nn.CrossEntropyLoss做了两件事先对模型输出做 Softmax 归一化成概率再计算交叉熵。为什么推荐直接调用因为它在数值稳定性上做了处理不会因为取对数时概率为0而产生 NaN。如果你希望完全不写循环里的损失表达式可以这样criterion nn.CrossEntropyLoss() loss criterion(outputs, labels)3.3 反向传播与优化器的工作流程训练循环核心代码一共四步optimizer.zero_grad() loss.backward() optimizer.step()很多人的疑问是为什么要zero_grad()。PyTorch 的backward()会把梯度累加到参数的.grad属性上如果不先清零上一轮的梯度会和本轮梯度叠加。虽然梯度累加在某些场景下是刻意为之的技术手段比如模拟更大的batch_size但在普通训练里是错误源。optimizer.step()根据梯度和优化算法更新参数。核心流程是前向传播 → 算损失 → 反向传播 → 更新参数 → 清零梯度循环往复。3.4 从零训练一个分类模型初始化相关参数如下学习率0.001到0.01之间SGD用0.01Adam用0.001。全过程代码示例model SimpleNN() criterion nn.CrossEntropyLoss() optimizer torch.optim.Adam(model.parameters(), lr0.001) epochs 10 for epoch in range(epochs): running_loss 0.0 for images, labels in train_loader: outputs model(images) loss criterion(outputs, labels) optimizer.zero_grad() loss.backward() optimizer.step() running_loss loss.item() print(fEpoch [{epoch1}/{epochs}], Loss: {running_loss/len(train_loader):.4f})loss.item()把单元素的张量转成Python数字。打印loss时如果发现剧烈波动不用急着调模型先降低学习率或增大batch_size。4. 训练与评估全流程实操理论知识说得差不多以后完整跑一个流程。这段实操记录里包含我反复用到的工具函数和判断逻辑。4.1 训练循环细节与进度监控上面核心循环已可见这里补两个实战函数。第一个是准确率计算def calculate_accuracy(outputs, labels): _, predicted torch.max(outputs, 1) total labels.size(0) correct (predicted labels).sum().item() return correct / totaltorch.max返回两个Tensor第一个是最大值第二个是最大值索引正好就是预测类别。第二个是控制训练时长的判断逻辑。当loss不再下降甚至反弹时常规的梯度下降已经优化不动了可能是学习率太大或模型容量不够。4.2 评估模式的切换与验证训练完成后测试集评估时有两件特别容易漏掉的事torch.no_grad()和model.eval()。def evaluate(model, test_loader): model.eval() correct 0 total 0 with torch.no_grad(): for images, labels in test_loader: outputs model(images) _, predicted torch.max(outputs, 1) total labels.size(0) correct (predicted labels).sum().item() accuracy 100.0 * correct / total return accuracymodel.eval()切换模型到推理模式Dropout层会关闭、BatchNorm层使用全局统计量。torch.no_grad()告诉PyTorch不要建反向传播图测试时不需要梯度省显存、提速。验证集和测试集在这里的区别是验证集用于每次epoch结束后的检查测试集只在训练全部结束后用一次。如果一直拿测试集调节超参数测试集就“污染”了最终分数会虚高失去无偏评估的意义。4.3 实验结果记录与分析用上述代码跑一轮通常在MNIST测试集上能到97%-98%的准确率。loss 收敛曲线长这样第一轮loss大约0.2-0.3后面逐渐下降到第六七轮接近0.05左右。如果训练集准确率接近100%测试集卡在97%不再上升主要表现为轻微过拟合可以尝试降低模型容量或增加正则化。有个很反直觉的现象是网络明明很大但没有任何过拟合——因为训练数据量很大6万张而网络参数并不多约12万模型实际参数量小于样本数过拟合压力相对有限。4.4 保存、加载和推理的完整路径训练完成后两步必做保存权重和单张推理。PyTorch保存模型有推荐方式torch.save(model.state_dict(), mnist_model.weights.pth)只保存权重state_dict而不是整个模型对象这样最稳妥。加载时先重建模型结构再load_state_dictmodel SimpleNN() model.load_state_dict(torch.load(mnist_model.weights.pth)) model.eval()推理流程与评估类似注意三个点灰度图单独一张记得扩展batch维image.unsqueeze(0)输入张量转float32不要用uint8输出结果在logits上取argmax得到类别编号5. 常见问题与排查技巧实录这段内容来自带新人和自己反复踩坑后的总结。前五类是初学者最容易遇到的。5.1 维度不对的报错RuntimeError: mat1 and mat2 shapes cannot be multiplied全连接层输入是784维。如果输入图像没展平或尺寸不对就会报这个错。排查入口是打印x.shape。如果图是3通道彩色图如32x32x3需要把3*32*32算进第一层输入input_size3072。5.2 dtype 不对Expected all tensors to be on the same devicePyTorch常见错误之一是类型不匹配例如标签张量是int64而模型输入要求float32。ToTensor()之后图片已经转float标签保持long类型即可。5.3 CUDA 设备不一致RuntimeError: Expected all tensors to be on the same device数据在CPU上、模型在GPU上的场景很典型。排查方法很简单打印model参数所在设备。把数据和模型统一放到同一个设备上的写法device torch.device(cuda if torch.cuda.is_available() else cpu) model.to(device) images, labels images.to(device), labels.to(device)5.4 测试时显存暴涨或报错最常见的原因是忘记torch.no_grad()导致推理时也构建了计算图。长时间推理需要这行代码原因体现在占用法上有梯度计算和没梯度计算完全是两个显存量级。5.5 loss 不下降或变成 NaN遇到loss变成NaN优先排查三类问题学习率过大输入数据有缺失值或inf标签有非法值常用应急操作是降低学习率或将lr0.01改成0.001测试再打印一下数据和标签范围。5.6 复现困难前面提到设置随机种子能解决大部分复现问题但以下因素也会让结果不可完全复现硬件差异、不同版本的CUDA/cuDNN、不同PyTorch版本的算子实现差异。只要你的代码在相同环境下能稳定复现就可以继续用。6. 扩展建议从“能跑”到“会用”不少朋友跑完MNIST后问下一步该做什么。我的建议是沿着下面几个方向做实验每个方向都能带来新的理解。6.1 换数据集验证网络结构的设计空间把MNIST换成Fashion-MNIST代码几乎不用改。效果更直观因为衣服鞋子的灰度图像在视觉结构上更复杂同样的全连接网络准确率会比数字分类低几个点。这能让你意识到数据特征分布对模型有直接影响。接着可以试试CIFAR-10。三通道彩色图直接喂给全连接网络效果很差同时也暴露了一个概念图像数据用view(batch, -1)展平喂全连接会丢掉空间结构信息。6.2 调整网络结构的几个方向增加一个隐藏层观察收敛速度和准确率变化把每层的神经元数从128换成256或512看参数量和训练时长怎么变在隐藏层后加Dropout观察验证集准确率是否提升、是否改善过拟合改网络结构有一个经验法则一次只改一个变量否则出问题不知道是哪个改动引起的。用表格记录每次实验的超参数和测试集准确率形成自己的对比基线。6.3 超参数与训练配置的搜索思路最简单有效的搜索方式是网格搜索在几个关键维度上各取几个候选值依次排列组合跑一遍。虽然耗时但规律一目了然。学习率可以从 {0.1, 0.01, 0.001} 选batch_size 从 {32, 64, 128} 选。跑完后会发现变化最敏感的参数不是网络宽度而是学习率。6.4 嵌入可视化工具的手段完全可以用几行代码记录loss变化然后画出曲线。更进阶的操作是用TensorBoard或Weights Biases做可视化监控。这类工具解决的问题很简单如果你想观察loss在第几轮开始过拟合靠打印日志很难直观感知画图会更清楚。7. 加快训练速度的操作训练时间长了以后很多人的诉求从“让准确率高一点”变成“同样结果跑快一点”。以下几点优化实际帮助很大。7.1 设备选择的优先级默认优先用CUDA。数据量小、模型小时CPU与GPU差距不明显但一旦数据量到几万张图片GPU优势立刻体现。如果没有独立显卡Google Colab的免费GPU环境可以满足入门需求。7.2 数据加载的并发优化train_loader DataLoader(dataset, batch_size64, shuffleTrue, num_workers2, pin_memoryTrue)num_workers控制子进程数量让数据加载与模型计算并行。Windows上建议先设置为0程序报错排查多个变量时尽量简化问题。pin_memory开启后能加速数据传输到GPU的过程。7.3 混合精度训练的入门认知PyTorch 2.x 提供了torch.amp混合精度接口用float16和float32混合计算显存占用能降低约一半训练速度在支持硬件上明显加快。入门阶段暂时不必深入学习细节但要知道有这条优化路径后续模型变大时可以回来补充知识。7.4 训练轮次的收敛判断与早停机制早停机制是实践中常用的训练策略每个epoch结束后用验证集计算准确率如果连续几个epoch验证集准确率不再提升就提前终止训练避免在过拟合区间浪费时间。实现要点是保存当前最优权重在早停时回滚到最优状态。对MNIST这类数据经验跑法一般10~15个epoch足够收敛再多对最终结果帮助不大。8. 真正要练的内容和坑位总结一个核心的训练理念可以提炼成这句话跑通是底线理解是全部分数。看一遍教程和自己动手写一遍区别很大。给别人讲一遍甚至把代码从零默写出来理解又完全不同。分享几个最容易被忽略的小技巧打印model对象能看到完整的网络结构和参数设置快速检查层定义是否有误打印type(images)、images.shape、labels.shape是排查维度问题的第一关每跑完一个epoch打印一次验证集准确率不要等到全部结束才看结果在长时间训练中断后恢复训练把optimizer.state_dict()一起保存更稳妥如果训练过程中遇到“为什么我的模型准确率一直很低”的问题先把数据标准化、learning rate、模型层数三个变量检查一遍这三个问题占了80%以上的原因。8.1 参数初始化为什么重要PyTorch的nn.Linear默认会使用某种自适应初始化方式所以大多数情况下代码会自动处理。但当你觉得自己改的模型收敛太慢时可以尝试手动设置初始化方式需要足够了解网络类型再操作。8.2 学习率调节的经验范围SGD一般从0.01开始Adam从0.001开始这两个值是行业里反复使用的经验起点。如果结果不理想稍微往下调低一个数量级试一下比来回乱试更高效。用一个较小数据集快速跑3~5轮对比loss走向是选择学习率的实用技巧。8.3 训练结果不稳定的源头排查如果两次相同配置训练最终精度差距超过1%先去检查数据加载是否设置shuffleTrue、随机种子是否固定。再看数据增强是否引入了随机因素。最后检查训练和推理之间是否用了同一个模型状态测试集评估前务必调用model.eval()。我自己的实践习惯是每轮实验前固定随机种子、记录每次实验的配置信息、训练结束立即保存最佳权重并测试。这三步坚持下来后期做任何实验迭代都省力很多。深度学习入门最怕的是“调好之后不知道为什么跑挂了也不知道为什么”。用这个思路跑通第一个分类网络后后面学卷积网络、Transformer、目标检测的路径会顺很多。把基本功练扎实再去搭更复杂的结构才会真正体会到里面各层组件的作用。