☰
用PyTorch从零构建第一个神经网络:MNIST实战与调试指南
2026/10/13 4:12:55 网站建设 项目流程

很多人在网上搜“用PyTorch构建你的第一个神经网络”,照着某个MNIST教程敲一遍,模型能跑了,准确率九成多,就以为自己入门了。但等到把图片从28x28换成32x32,把数据集换成自己的业务数据,马上满头问号:为什么我的模型不收敛?为什么batch size设大了内存爆炸?为什么同样的代码换台电脑效果差那么多?

问题出在“照抄demo”而不是“理解机制”。所以我这篇不是给你一份可以复制粘贴的完整代码,而是带着你把一个最简单的神经网络从零抠明白:每一行代码在干什么,每个参数为什么是那个值,报错时先从哪个方向查。你跟着走完一遍,再看任何其他网络结构,都不会觉得那是个黑盒子了。

我默认你已经知道Python的基本语法,装好了PyTorch,但没怎么亲手训过模型。这篇文章的目标很直接:让你独立完成“加载数据 — 定义网络 — 训练 — 评估”的完整闭环,并且在这条路上避开我当年踩过的坑。

1. 选MNIST做第一个项目的理由:玩具任务才能逼出真问题

1.1 为什么不是CIFAR,也不是你自己的数据

几乎所有入门教程都会选MNIST手写数字,很多人嫌它“太简单、没意思”。但我的真实体会是:第一次训练神经网络,你需要的不是挑战,而是可控性。MNIST单张图片只有28x28像素,展开成全连接层的输入也就是784个数字,模型可以在普通CPU上几分钟内跑完一个epoch。这意味着你做任何改动,都能很快看到结果反馈,调试反馈越短,学习效率越高。

CIFAR-10虽然也只有10个类别,但它是32x32的彩色图片,输入维度变高,训练时你还要考虑数据增强、归一化方式、可能的过拟合问题。这些对于第一次做的人来说全是干扰项。你自己的业务数据就更不用说了,数据清洗和标签问题会让你根本分不清是模型的问题还是数据的问题。

MNIST这个任务还有一个好处:它有公认的标准答案。测试集准确率能到97%~98%是基本盘,如果你的模型连95%都到不了,说明代码里一定有bug。这种“可验证性”是新手阶段最宝贵的——你知道问题出在代码而不是玄学。

提示:第一直觉不要追求模型多复杂、准确率多高。先把一条最简单的链路跑通,再逐步加东西。新手最常见的错误是环境装完直接从GitHub拉一个ResNet的demo开始跑,报错以后连是环境问题还是代码问题都分不清楚。

1.2 环境准备的“最小可行版本”

我建议你直接在本地装PyTorch的CPU版本起步,不需要一开始就惦记CUDA。GPU能加速训练,但也意味着你还要学会处理CUDA版本和PyTorch版本匹配的问题,这在学习阶段完全没有必要。等你在CPU上把那套送代逻辑彻底弄懂了,再上GPU就是改一行设备代码的事。

安装就按官方命令来,根据自己的操作系统选对应版本。我在工作机上是这么装的:

pip install torch torchvision

装完验证一下能否正常导入:

python -c "import torch; print(torch.__version__)"

能输出版本号,环境就基本就绪了。之后在Jupyter Notebook或者任意Python编辑器里写代码都行。我自己习惯用Jupyter,因为中途想看看某个张量的shape、打印几个结果都很方便,调试体验比在命令行里写整脚本友好得多。

注意:torchvision里带的是MNIST数据集,你不需要自己准备图片。但如果你身处某个网络环境,第一次下载数据可能要等一会,这很正常。下载失败就看下是否能手动下载数据文件放到本地目录,入门阶段不值得在环境上折腾太久。

2. 神经网络到底在学什么:从线性变换到非线性叠加

2.1 一张图片怎么变成一堆数字

新手最容易懵的地方是“模型输入输出到底是什么”。拿MNIST来说,一张图片在torchvision里读进来,是一个形状为(1, 28, 28)的张量,其中“1”是颜色通道(灰度图只有1个通道),高和宽都是28。把每个像素看作一个数值,黑色接近0,白色接近255,这张图本质上就是784个0到255之间的数字。

我们要做的,是让模型学会:给定这784个数字,输出一个10维的向量,这10个维度分别代表“是0的概率”“是1的概率”……“是9的概率”。模型内部怎么做到这件事?第一层做的事情其实非常简单——每个输出节点都和784个输入做加权求和,再加上一个偏置。

假设第一层有128个节点,那就是128组权重,每组784个数字。你有多少权重,模型就有多少“参数”。参数一开始是随机初始化的,也就是说模型最初完全是在瞎猜。所谓训练,就是反复调整这些权重,让模型的输出逐渐逼近正确答案。这就是神经网络的全部本质:调整一堆数字,让映射关系从瞎猜变成可靠。

2.2 非线性激活函数为什么不能省

如果只做矩阵乘法,那不管堆多少层,本质上都等价于一层线性变换,模型能力天花板就是一条直线。MNIST的像素和数字类别之间的关系显然不是一条直线能画出来的,所以你需要在线性层之间插入非线性函数。

我第一个项目用的是ReLU(Rectified Linear Unit)。它的数学形式很简单——小于0的部分压成0,大于0的部分原样保留。这个操作有什么用?ReLU让每一层输出的某些节点“熄灭”、某些节点“激活”,整个网络就有了近似任意复杂函数的能力。

用大白话说:线性层负责画很多条直线,ReLU负责在合适位置“折断”这些直线,让它们拼接成曲线。层数越多,能拼出的形状就越复杂。这也是为什么只加一层Linear往往效果不够,因为你没有给网络足够的“折叠”机会。

提示:新手可以先完全照搬ReLU这个选择,不要纠结为什么不用sigmoid或tanh。这个阶段你要理解的是“非线性必须有,否则网络白搭”,等基本功扎实了再去研究不同激活函数的梯度特征和适用场景。

3. 第一个模型骨架:用Sequential组装线性层与激活函数

3.1 一个三层的全连接网络

用PyTorch定义网络最简单的方式是torch.nn.Sequential,它像搭乐高一样把层按顺序拼起来。下面是新手阶段最合适的结构:两层隐藏层加一个输出层,中间各夹一个ReLU。

import torch.nn as nn class MyFirstNet(nn.Module): def __init__(self): super(MyFirstNet, self).__init__() self.layers = nn.Sequential( nn.Linear(784, 128), nn.ReLU(), nn.Linear(128, 64), nn.ReLU(), nn.Linear(64, 10) ) def forward(self, x): x = x.view(x.size(0), -1) return self.layers(x)

这段代码里最值得新人抠细节的是第一个Linear的参数,784和128分别是什么。784是输入维度,因为是28x28展平后的长度;128是我指定的隐藏层宽度,即“这一层要提取多少个特征”,这个数字本身是超参数,没有唯一答案。调大一点模型容量更大但也更容易过拟合,调小一点更快但可能学不到位。新手阶段先照抄128,后续可以自己改成256对比一下。

第二个Linear是128到64,这代表第二层把第一层提取出的128维特征压缩成64维。最后一层是64到10,10就是类别数,输出10个数,后续会转成10个概率。注意最后一层没有接ReLU,原因后面讲损失函数时会提到。

3.2 forward里的view是干什么的

你在很多教程里会看到x = x.view(x.size(0), -1)这行,新手经常忽略它。这是因为全连接层nn.Linear要求的输入是二维形状(批次大小, 输入维度),而数据从DataLoader里出来时是(批次大小, 1, 28, 28),多了通道和二维空间结构。view操作就是把后面三个维度压扁成一个784维的向量。

其实你也可以在把数据喂给模型前就做好展平。我习惯在forward里做,好处是模型内部对外界输入格式的容错性更高,后面如果换数据集、换输入尺寸,只需要改view逻辑。

print(x.shape) # 比如 torch.Size([64, 1, 28, 28]) x = x.view(64, -1) print(x.shape) # torch.Size([64, 784])

这里-1是自动推导,PyTorch看到你有64张图、总元素数是64128*28,就会自动算成每张图784。

4. 训练循环拆解:从loss到梯度再到权重更新

4.1 训练循环里的五行核心代码

模型定义完成之后,真正让网络“学会”的是一个循环。每个深度学习框架的训练循环长得都差不多,核心就这几步:算预测值、算损失、清零梯度、反向传播、更新参数。

import torch.optim as optim criterion = nn.CrossEntropyLoss() optimizer = optim.Adam(model.parameters(), lr=1e-3) def train_one_epoch(model, train_loader, criterion, optimizer): model.train() total_loss = 0 for images, labels in train_loader: outputs = model(images) loss = criterion(outputs, labels) optimizer.zero_grad() loss.backward() optimizer.step() total_loss += loss.item() return total_loss / len(train_loader)

我逐个说。这里的outputs是模型对一整批图片的输出,形状为(批次大小, 10),表示每个样本属于各类别的原始得分。loss是模型得分和真实标签之间的差距,PyTorch的交叉熵函数会自动处理两个事情:把得分变成概率,然后计算和真实标签之间的损失。这也是模型最后一层不加ReLU的原因——CrossEntropyLoss内部包含softmax操作,你在网络里加了ReLU反而会破坏输出的分布性质。

optimizer.zero_grad()必须放在每次更新前。PyTorch的梯度是累加的,不手动清零的话,第二次反向传播会把第一次的梯度叠加进去,导致梯度变成一串乱账。这是新手最容易漏的细节,漏掉的结果就是loss曲线乱跳或者完全不能收敛。

4.2 backward和step之间隔着一个梯度下降的“慢思考”

loss.backward()做的事是计算网络里每个参数对loss的梯度,也就是“每个权重调大一点会让loss变大还是变小、变多少”。这一步是整个训练过程的灵魂。你可以直观理解为:loss是当前错误程度,梯度是让错误下降的方向导引。

optimizer.step()做的事才是真正“往前走一步”。它根据刚算好的梯度和算法设定,把每个参数更新一下。对于SGD来说就是param = param - lr * gradient,Adam稍微复杂一点,它会给不同参数自适应不同的学习率。这也是我推荐新手直接用Adam的原因——它对学习率的敏感度比SGD低得多,少一个难调的旋钮,前期省心很多。

你会发现整个循环里没有一行代码在“拿来验证”。所以每次epoch结束还需要单独跑一个评估循环,把模型切成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() return correct / total

这里为什么要有with torch.no_grad()?因为评估时我们只想看看模型表现,不需要计算梯度。如果不开这个,系统还是会默默构建整个计算图、记录中间变量的历史,白白浪费内存和算力,数据量大时甚至会直接内存不足。

torch.max(outputs, 1)的意思是沿着每行取最大值,返回(最大值, 对应下标),下标就是模型预测的类别。这一行在各类分类任务里几乎是标配写法,建议直接背下来。

5. 你大概率会踩的坑:从数据维度到loss不降的完整排查链路

5.1 维度不匹配:最常见但也最好解决

跑第一个训练循环时,很多人第一眼看到的报错是类似“Expected input batch_size (64) to match target batch_size (32)”,这说明数据在某个环节丢了样本,或者标签和图片没有对齐。这通常是DataLoader里设置了几个worker导致的,或者是数据增强环节里有的操作返回了空样本。入门阶段建议直接用默认的num_workers=0,不要为了速度盲目开多进程,排查起来会简单很多。

另一个维度问题是:你看了云端示例代码,示例里的模型第一层直接是nn.Linear(3, 32),因为对方用的是三通道彩色图。你拿自己的单通道灰度图一跑,直接报shape不匹配。这不是PyTorch版本问题,而是输入通道数对不上模型第一层的in_features。遇到这种报错,先检查模型第一层的输入维度和数据实际形状,不要着急搜报错内容。

5.2 loss不降或者震荡的排查顺序

如果你训练了十几个epoch,loss纹丝不动,或者先是剧烈震荡然后一路狂奔,下面是按优先级排列的排查思路,照着这个顺序来,不要跳:

第一,确认数据归一化做了没有。MNIST如果不除以255,像素值在0~255之间,梯度尺度会非常不稳定,优化器很容易在早期把权重更新的乱七八糟。入门代码里transforms.ToTensor()已经顺手把像素变成了0到1之间的浮点数,如果你用的是自己整理的数据,一定要做类似的归一化。

第二,检查学习率是不是大了一两个数量级。我见过太多新手在Adam下用0.1或0.01,结果loss曲线像心电图一样乱跳。1e-3这个值对于这种小网络是安全的起点。如果loss一开始就变成nan了,那大概率是学习率太大导致梯度爆炸。降到1e-4或者1e-5跑两步再看,基本都能恢复正常。

第三,确认训练循环里是否忘了常见的三个操作:model.train()、model.eval()、optimizer.zero_grad()。我把它们称为“新手三必忘”。忘记zero_grad的后果在2.1节说过了,模型跑起来准确率很低而且怎么调都没用。

第四,如果上面都查完还是不对,去看数据加载部分。直接在训练循环开头print一下每张image的shape和labels的shape,再用torch.unique(labels)确认标签真的是0到9,而不是什么文件读错了、标签全变成了同一个值。

5.3 准确率能收敛但始终在90%上不去的后续优化

如果你已经能稳定收到90%以上的准确率,但卡在某个值上不去,恭喜你,这说明代码链路是通的,你进入超参数调试阶段了。第一个值得做的是加大隐藏层宽度,比如把128改成256,一般能立刻看到提升,代价是训练时间变多。其次是尝试把batch size从64调到128,有时候小batch反而梯度噪声大,大batch会更稳定。

再往下可以调epoch数。我发现很多新手训练只跑5个epoch就下结论说“效果不好”,其实小网络在这个任务上跑到20~30个epoch仍可能持续变好。你可以在记录每个epoch的准确率,看曲线是否还在上升,如果还在涨就多跑几个epoch。这个阶段耐心比技术有用。

提示:我们第一次用PyTorch的时候,记得当时上传了一坨代码到某个开源社区,被一个陌生人回复了一句话:“你的loss曲线在第十个epoch之后还在降,你只跑了15个epoch,为什么就不跑了?”那句话比任何教程都让我长记性——先让训练跑完,再判断问题。

6. 让模型真正能投入使用:评估之外的几个细节

6.1 保存和加载模型这件事也有讲究

训练完模型,你当然希望把它保存下来,下次不用重新训练。PyTorch里保存模型的方式有两种。完整模型直接torch.save(model, 'model.pth'),这种做法把网络结构和权重一起存了。我建议你用state_dict的方式保存,也就是只存参数字典:

torch.save(model.state_dict(), 'mnist_model.pth') # 恢复 model = MyFirstNet() model.load_state_dict(torch.load('mnist_model.pth'))

原因很朴素:state_dict的保存方式只存数字,不牵扯Python对象序列化,换环境、改代码、跨PyTorch小版本升级时问题少得多。complete model的保存方式在加载时会绑定当时的类定义,你稍微改一点模型结构就加载不上了。养成存state_dict的习惯,以后在项目里也用得到。

如果你打算之后换GPU推理,加载时要多传一个map_location参数,把权重映射到CPU或指定设备上,不然报错也是家常便饭。

6.2 从单张图片测试到观察模型内部

跑完测试集准确率,建议你额外做一件事:拿手边随手画的数字测试模型。你可以用PIL画一个数字,然后预处理成和训练数据一样的格式——灰度、28x28、像素值范围0到1、转成Tensor——再喂给模型,打印每个类别的概率。这会让我第一次有“这玩意儿真的在工作”的感觉。

输出的10个数字里很可能不是一个干干净净的概率分布,比如某个样本类别7的得分是8.2,类别1的得分是7.9,两者非常接近。这个信息很有价值。你可以写一小段代码找出来哪些样本模型预测得很犹豫,然后挑出来看一眼。这些边界样本就是你下一步优化的方向:是数据里这个写法比较罕见,还是模型容量不足,信息量大得很。

7. 下一步的多种玩法:从MNIST毕业之后怎么进阶

如果你踏踏实实把这篇文章里的代码全部手敲一遍,并且耐着性子调过一轮超参数,那你的下一步每条路都是通的。我个人比较推荐的反而不是立刻冲去做图像分类比赛,而是顺着两个方向走:一是换一个结构化数据集,体会数据预处理和特征工程的重要性;二是把全连接网络换成卷积网络,观察模型是如何通过卷积核自动提取局部特征的。二选一都能让你对“深度学习是怎么工作的”多一层实感。

换网络结构这件事,本质上还是同一套训练循环:定义模型、选损失函数、创建优化器、循环训练、保存权重。你会发现训练代码一行都不用改,只改模型定义部分。这恰恰说明你第一个神经网络项目学得有多扎实。

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

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

立即咨询