PyTorch从零构建MNIST手写数字识别：完整训练流程与调试心得

1. 项目概述：第一个神经网络该从哪下手

第一次用 PyTorch 搭神经网络，很多人会卡在同一个地方：教程看了不少，但代码在自己电脑上跑不起来。这个项目的目的很纯粹，用 PyTorch 从零构建一个能完成手写数字识别（MNIST）的神经网络，把“训练一个模型”这件事的完整链路跑通。只有亲手跑通一次，后续学卷积神经网络、循环神经网络、Transformer 这些结构时才不会心虚。

为什么拿 MNIST 而不是更高难度的数据集？因为 MNIST 足够简单，图片是 28×28 的灰度图，类别只有 0 到 9。模型结构哪怕非常朴素，也能轻松达到 97% 以上的准确率。这意味着调试成本极低，新手可以把精力放在理解 PyTorch 本身的工作方式上，而不是和模型不收敛、数据太复杂这类问题纠缠。

适合谁来参考？你只要会用 Python，了解最基本的 NumPy 操作，甚至没写过神经网络都可以。我会从环境搭建讲起，把数据加载、模型定义、训练循环、验证与保存这几个环节全部拆开，给出可以直接运行的代码。所有代码基于 PyTorch 2.x，但 1.x 也完全兼容。读完你不仅能跑通这个项目，还能知道每一步在做什么、为什么要这么做。

2. 环境准备：PyTorch 安装与项目结构

2.1 安装 PyTorch 时最常见的坑

安装 PyTorch 看起来是一条命令的事，但“用哪个命令”很讲究。PyTorch 官方首页会给你一个根据操作系统、包管理工具、CUDA 版本自动生成的安装命令，但这个命令有个大坑：如果直接复制默认命令，大概率装的是 CPU 版本，或者装了一个和你显卡驱动不匹配的 CUDA 版本。

先判断自己需要什么版本。如果电脑有 NVIDIA 显卡，可以在命令行里输入 nvidia-smi 查看驱动支持的 CUDA 版本，右上角会显示类似 CUDA Version: 12.1 的字样。注意，这个版本号是驱动能支持的最高版本，安装 PyTorch 时选的 CUDA 版本只要不高于它就行。比如驱动支持 12.1，就可以放心选 cu121 或者 cu118 的 PyTorch。

安装命令通常长这样：

bash复制# CPU 版本
pip install torch torchvision torchaudio

# CUDA 11.8 版本
pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118

# CUDA 12.1 版本
pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu121

新手最容易犯的错，是先装了一个 CPU 版，后面发现训练太慢，又去重装。实际上同一环境里重装 PyTorch 很容易出现依赖冲突，比如缺 nvidia-cuda-runtime 之类的底层库。我建议在动手之前先花两分钟用 nvidia-smi 确认一下，而且强烈建议创建独立的虚拟环境，不要直接装在系统 Python 里。

bash复制conda create -n pytorch_env python=3.10
conda activate pytorch_env

如果你用的是 Anaconda 或者 Miniconda，这一步能避免以后 90% 的环境混乱问题。不用的项目各自独立环境，版本互不干扰，出了问题直接删掉重建，比费尽心思去修复依赖关系痛快得多。

2.2 验证安装是否成功

安装完成后，千万别说“装完了就开始写代码”。先做一个五秒钟的冒烟测试：

python复制import torch
print(torch.__version__)
print(torch.cuda.is_available())

如果第一行输出版本号，说明 PyTorch 本体装好了。第二行如果输出 True，说明 CUDA 可用；如果输出 False，说明你装的是 CPU 版或 CUDA 配置有问题。想确认 GPU 是否真的参与计算，可以再加两行：

python复制print(torch.cuda.get_device_name(0))
x = torch.randn(3, 3).cuda()
print(x)

只要看到能正常输出 GPU 的名字和一个 3×3 的张量，环境就彻底没问题了。这一步别省，后面所有问题排查都要基于一个“确定可用”的环境，否则训练时遇到报错会分不清是代码问题还是环境问题。

3. 数据准备：MNIST 数据集与 DataLoader

3.1 下载数据集时的网络问题与处理

MNIST 数据集在 PyTorch 里可以直接下载，不需要手动去官网找文件。标准写法是：

python复制from torchvision import datasets, transforms

transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.1307,), (0.3081,))
])

train_dataset = datasets.MNIST(root='./data', train=True, download=True, transform=transform)
test_dataset = datasets.MNIST(root='./data', train=False, download=True, transform=transform)

这里有两个细节值得注意。第一个，transforms.ToTensor() 会把 PIL 图片或 NumPy 数组转换成张量，并且自动把像素值从 0 到 255 缩放到 0 到 1。第二个，Normalize((0.1307,), (0.3081,)) 是 MNIST 数据集的全局均值和标准差。这两个数值是官方统计好的，直接用就行。

归一化这一步很多人不理解为什么要做。简单说，输入数据的量级如果不统一，网络训练时会非常不稳定，梯度下降过程会像在高低不平的地面上跑来跑去。把数据缩放到均值为 0、标准差为 1 的分布，相当于把地面整平了，模型收敛速度快非常多。

下载慢是常见问题。PyTorch 从境外服务器下载 MNIST 有时会很慢或直接超时，尤其是在国内网络环境下。遇到这个问题，可以考虑用 gitee 或国内镜像站的数据源，或者手动下载四个 .gz 文件放到 ./data/MNIST/raw/ 目录下，让代码检测到文件已存在就不会再去下载。手动放置时文件命名必须和源码里一致，分别是 train-images-idx3-ubyte.gz、train-labels-idx1-ubyte.gz、t10k-images-idx3-ubyte.gz、t10k-labels-idx1-ubyte.gz。

3.2 DataLoader 到底在做什么

数据加载这一步，初学者容易直接把整个数据集一次性扔进模型。MNIST 只有 6 万张训练图片，每张是 1×28×28 的张量，全部装进内存也才一百多 MB，看起来好像没问题。但这样做的坏处在于：模型参数更新需要随机梯度下降，一次性把全部数据算完再更新，既不随机，显存也不够用。

DataLoader 的作用是把数据集切成小批量（batch），每个 batch 送入模型训练一次。它内部还负责打乱数据顺序（shuffle），避免模型因为相似数据连续出现而产生偏见。shuffle 只在训练集上用，测试集不需要，因为测试只是纯粹的前向计算，不打乱还能稳定复现结果。

python复制from torch.utils.data import DataLoader

batch_size = 64
train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False)

batch size 设多少合适？这个超参数直接影响训练速度、显存占用和模型收敛质量。64 在 MNIST 上是一个合理的选择。太小会让梯度估计的方差大，训练不稳；太大会占显存，而且可能导致收敛到泛化性能较差的极值点。具体的权衡后面在调试部分再展开。

在写训练循环之前，可以先随便取一个 batch 看看数据的形状，确认自己的理解正确：

python复制data, target = next(iter(train_loader))
print(data.shape)   # 应该是 torch.Size([64, 1, 28, 28])
print(target.shape) # 应该是 torch.Size([64])

data 的四个维度分别是 batch 大小、通道数、高度、宽度。MNIST 是灰度图，所以通道数是 1；如果是彩色图就是 3。这种形状约定在你后面用卷积神经网络时会反复遇到，一定要记住。

4. 模型定义：前馈神经网络的结构与实现

4.1 用 nn.Module 定义自己的网络

PyTorch 中定义模型的标准做法是继承 nn.Module 类。虽然也可以直接用 nn.Sequential 快速堆层，但继承类的方式更灵活，逻辑也更清晰，对新手理解模型结构帮助更大。

第一版可以设计一个三层的全连接网络。输入层 784 个神经元，对应 28×28 的图片展开成一维向量；两个隐藏层分别 256 和 128 个神经元；输出层 10 个神经元，对应数字 0 到 9。代码：

python复制import torch
import torch.nn as nn
import torch.nn.functional as F

class SimpleNN(nn.Module):
    def __init__(self):
        super(SimpleNN, self).__init__()
        self.fc1 = nn.Linear(28 * 28, 256)
        self.fc2 = nn.Linear(256, 128)
        self.fc3 = nn.Linear(128, 10)

    def forward(self, x):
        x = x.view(x.size(0), -1)    # 把 1×28×28 展平成 784
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = self.fc3(x)
        return x

这里的关键点是 forward 方法。你可能听说过 PyTorch 的“动态图”机制，本质上就是：每次前向计算时，代码是从上往下一条一条执行的，每一步都记录计算图。这和 TensorFlow 1.x 时代先建图再跑会话完全不同，调试起来直观得多。

展平操作 x.view(x.size(0), -1) 是个高频操作，意思是保留 batch 维度，其余所有维度展平成 1 维。-1 会自动推导，所以 64×1×28×28 展平后是 64×784。如果不用 view 而直接用 x.reshape 也可以，但 view 在大多数情况下更高效，因为它尽量不复制数据。

4.2 激活函数为什么必须加

如果你把代码里的 F.relu 删掉，模型就变成了三个线性层的简单堆叠。你可能会以为层数越多模型越强，但数学上有一个令人失望的事实：多个线性层的组合本质上还是一个线性变换。也就是说，没有激活函数的深层网络和一个单层网络表达能力完全相同。

ReLU 函数本身极其简单：输入小于 0 时输出 0，输入大于等于 0 时输出原值。但正是这个简单的非线性操作，让网络可以逼近任意复杂的函数。放在两个层之间，相当于给每一层都赋予了“决定哪些信息可以通过”的能力。

选择 ReLU 而不是 Sigmoid，主要是为了解决梯度消失问题。Sigmoid 函数在输入绝对值很大时梯度趋近于 0，在深层网络中梯度一层层乘过去，很快就消失了，前面的层根本得不到更新。ReLU 在正半轴的梯度恒为 1，有效缓解了这个问题。实际项目里 ReLU 已经是默认首选。

还有一个细节：输出层不能加激活函数。如果是分类任务，模型输出的 10 个原始数值叫 logits，要经过 softmax 才能变成概率。但 PyTorch 的损失函数 CrossEntropyLoss 内部已经包含了 softmax 操作，所以模型输出层直接给 logits 就行，不需要自己再包一层 softmax。新手常犯的错就是画蛇添足，在输出层额外加一个 F.softmax，反而会导致训练效果变差。

4.3 通过参数量理解网络规模

写完模型后值得做一件事：算一下这个简单网络有多少参数。

python复制model = SimpleNN()
total = sum(p.numel() for p in model.parameters())
print(total)

第一层是 784×256 的权重，加上 256 个偏置，一共约 20.1 万参数。第二层是 256×128，约 3.3 万。第三层是 128×10，约 1290。总计约 23.5 万参数。这些参数就是模型需要学习的全部内容。

这个数字看起来不少，但和现代大模型动辄百亿千亿参数相比，根本不值一提。理解这一点有助于建立一种感觉：模型容量的大小由参数量决定，而选择合适的容量是工程中的核心问题。MNIST 用 20 万参数已经非常充裕了，再加大网络收益有限，反而更容易过拟合。

5. 训练循环：从损失函数到参数更新

5.1 损失函数与优化器的选择

分类任务的标准损失函数是交叉熵。PyTorch 里直接用 nn.CrossEntropyLoss()，它的内部计算是把模型输出的 logits 先做 softmax，再计算交叉熵。定义优化器时，最常用的也最稳妥的是 Adam：

python复制device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = SimpleNN().to(device)
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)

这里有一个关键的选择逻辑。传统的随机梯度下降（SGD）需要手动调节学习率、动量等超参数，对新手不友好；Adam 可以看作加了“自适应学习率”的梯度下降，每个参数都有自己独立的学习率，在一大批问题上都表现稳定。MNIST 这种简单任务，Adam 几乎不会出问题。

学习率 0.001 也是 Adam 最常用的默认值。学习率太大，loss 会震荡甚至爆炸；学习率太小，收敛极慢。后续你可以尝试把学习率调成 0.01 或 0.0001 看看差别，这是体会超参数影响最快的方式。

5.2 训练循环的每一行代码都要理解

PyTorch 里一个典型的训练循环长这样：

python复制epochs = 5

for epoch in range(epochs):
    model.train()
    running_loss = 0.0

    for data, target in train_loader:
        data, target = data.to(device), target.to(device)

        optimizer.zero_grad()

        output = model(data)
        loss = criterion(output, target)

        loss.backward()
        optimizer.step()

        running_loss += loss.item()

    avg_loss = running_loss / len(train_loader)
    print(f"Epoch {epoch+1}/{epochs}, Loss: {avg_loss:.4f}")

逐行解释：

model.train() 是告诉模型现在是训练模式。PyTorch 里有 train() 和 eval() 两种模式，对全连接网络影响不大，但当你使用 Dropout 或 BatchNorm 时，这两种模式的行为完全不同。新手可能在这里埋下隐患——训练时忘了写 model.train() 或者测试时忘了写 model.eval()，结果模型结果特别奇怪。

optimizer.zero_grad() 的作用是清空上一轮迭代累积的梯度。这一步非常关键。PyTorch 的梯度是累积的，不清空的话，新的梯度会加到旧梯度上，导致参数更新方向完全错误。很多初学调试半天找不出 bug，最后发现只是忘了做这一步。

loss.backward() 是反向传播，计算所有参数的梯度。PyTorch 的自动求导机制在这时根据计算图反向逐层求出损失对每个参数的偏导数，并把值存到每个参数的 .grad 属性里。

optimizer.step() 是根据优化器公式更新参数。注意 zero_grad 必须在 backward 之前放，因为一旦 backward 算好了梯度，下一步更新后梯度就失去意义，必须清零等下一个 batch 重新计算。

5.3 观察 loss 下降是判断模型是否学习的核心手段

训练过程中最直接的反馈就是 loss。第一轮结束时 loss 应该在 0.2 到 0.4 之间，如果正常的话，5 个 epoch 后可以降到 0.03 左右。loss 下降速度变缓是正常的，因为后期梯度变小，模型只是在做精调。

我建议你在第一次训练时把每 100 个 batch 的 loss 打印出来，而不是只打印每个 epoch 结束后的均值：

python复制if batch_idx % 100 == 0:
    print(f"Epoch {epoch+1} Batch {batch_idx} Loss: {loss.item():.4f}")

这样你能看到 loss 在单个 epoch 内部的波动规律。正常情况下它是震荡下降的，因为不同 batch 的数据难度不同，loss 一定会有起伏。如果 loss 完全不动或者反而上升，那就要停下来排查问题了。

训练完成后，不要急着说“模型训练完了”。我们要用测试集验证模型在没见过的数据上的表现。很多新手只看训练集准确率，发现 99% 就觉得成功了，结果部署到真实数据上效果一塌糊涂。这就是过拟合，模型把训练集的特征背下来了，却没学会泛化规律。

python复制model.eval()
correct = 0
total = 0

with torch.no_grad():
    for data, target in test_loader:
        data, target = data.to(device), target.to(device)
        output = model(data)
        pred = output.argmax(dim=1)
        total += target.size(0)
        correct += (pred == target).sum().item()

print(f"Test Accuracy: {100.0 * correct / total:.2f}%")

torch.no_grad() 是测试时的标准姿势。这个上下文管理器告诉 PyTorch 不需要记录计算图，也不需要保存中间变量，因此可以大幅降低显存占用和计算时间。测试阶段只做前向传播，用 argmax(dim=1) 取每个样本概率最大的类别作为预测结果。

MNIST 全连接网络 5 个 epoch 一般就能到 97% 左右的准确率。如果你观察到这个数字，恭喜你，第一个神经网络已经完整跑通了。

5.4 模型保存与加载

训练好的模型必须能保存下来，否则关掉程序就要重新训练。PyTorch 里推荐的做法是保存模型的 state_dict，也就是所有参数值，而不是直接保存整个模型对象。

python复制torch.save(model.state_dict(), "mnist_simple_nn.pth")

# 加载时
model = SimpleNN()
model.load_state_dict(torch.load("mnist_simple_nn.pth", weights_only=True))
model.eval()

有一点要特别注意：加载模型前必须先定义好和原来完全相同的模型结构。因为 state_dict 只是参数数值，它不知道这些参数属于哪个层。如果结构不一致，加载时会报 key 不匹配的错误。

继续使用这个模型进行预测时，输入数据同样需要经过预处理流程。因为训练时用的数据是归一化后的标准分布，测试时如果直接输入原始像素值，模型会认为数据分布变了，预测质量就会下降。这是新手很容易忽略的细节。

6. 调试心得：训练路上的典型问题与排查方法

6.1 数据和设备不匹配

在 CPU 上训练完的模型参数都在 CPU 内存里，如果 GPU 可用，最好把模型和数据都搬到 GPU。新手最常见的报错之一是“Expected all tensors to be on the same device”，这表示有些数据在 CPU、有些在 GPU。对策很简单：定义好 device 后统一调用 .to(device)，模型在初始化后调用一次，每个 batch 的数据在循环里调用一次。把这一步做在前面，后面会少很多麻烦。

6.2 loss 变成 NaN 或 inf

这个问题在训练中遇到会很崩溃。常见原因有三个：学习率过大，梯度爆炸，或者数据存在缺失值。MNIST 这个例子基本不会出现 NaN，但如果你自信地把学习率调到 0.1 以上，很可能会看到 loss 变成 nan。

排查顺序：先降低学习率，如从 0.001 降到 0.0001，看问题是否消失；再检查输入数据是否包含异常值或缺失值；最后考虑梯度裁剪（gradient clipping）。Adam 本身比 SGD 稳定，但如果损失函数或数据有问题，它一样救不了。

6.3 训练准确率低，loss 一直不降

这种情况很可能是下面几个原因：

• 归一化被遗漏或参数错误。很多人下载完数据后没做 ToTensor 和 Normalize，直接把 0 到 255 的像素值送给模型。这个情况下模型也能训练，但收敛非常慢，准确率上限也低。
• 激活函数缺失。前面说过，没有激活函数，多层线性层等于一层。
• 模型容量不够。MNIST 用 20 万参数足够了，但如果你只留一个 10 个神经元的隐藏层，模型表达能力太弱，准确率很难超过 85%。
• 数据加载乱序。训练时 shuffle=False 可能让数据按标签顺序排列，模型容易陷入偏见。确认 DataLoader 中 shuffle=True。

6.4 GPU 利用率太低

在 MNIST 这种小数据集上训练，如果你开 GPU 监控（比如 nvidia-smi dmon 或者更直观的图形监控工具），可能会发现 GPU 利用率只有 30% 甚至更低。这是正常现象。大多数计算发生在 CPU 处理数据加载、放大到 GPU 的环节，GPU 自己真正计算的时间很短，无法吃满。

但这也能反映一个优化思路：如果以后用更大的数据集，数据加载可能会成为训练速度的瓶颈。可以在这个阶段就养成一个好习惯——使用 DataLoader 的 num_workers 参数：

python复制train_loader = DataLoader(train_dataset, batch_size=batch_size,
                          shuffle=True, num_workers=4, pin_memory=True)

num_workers 可以启动多个子进程做数据预处理，pin_memory=True 则把数据放进锁页内存，让数据传输到 GPU 更快。这两个参数的组合在 Windows 上偶尔会报多进程错误，加上 if __name__ == "__main__": 包裹训练代码就能解决。

6.5 复现性的基础设置

如果希望实验结果可以被固定下来，比如今天跑出的准确率和明天跑出的完全一致，需要固定随机种子。PyTorch 和 Python、NumPy 各有一套随机数生成机制，都要设置：

python复制import random
import numpy as np

def set_seed(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

在训练前调用 set_seed(42) 即可。注意，即使做了这些，GPU 浮点运算的并行累加顺序也可能带来极小的误差，但绝大多数情况下已经够用了。

7. 经验总结：从第一个网络到后续进阶

第一版全连接网络跑通之后，你已经有了一条完整的技术链路：数据准备 → 模型定义 → 训练循环 → 验证评估 → 保存加载。这个骨架在所有 PyTorch 项目里都一样，后面学卷积神经网络只是替换模型结构，学循环神经网络只是改变数据输入维度，学 Transformer 也是在这个框架上增加更复杂的模块。

我再分享一个自己在实际操作中拜访过很多次的小技巧。训练时如果总感觉模型效果不好，不要只盯着 Accuracy 一个指标。先把测试集中被预测错误的样本可视化出来，用 matplotlib 画成网格图，同时打印出模型预测的类别和真实类别。你会异常清晰地发现错误集中在哪些数字之间，比如 4 和 9 这类形状相近的数字。有了这种直觉后，再去调整模型或改进数据增强，就是有方向地做事，而不是靠蒙。

对刚接触 PyTorch 的读者，我建议在当前这个项目的基础上尝试几个小改动：

• 把优化器从 Adam 换成 SGD 并调高学习率到 0.01，体会训练过程的差异；
• 增加一个隐藏层或者把每层的神经元数量翻倍，观察性能和训练时间的变化；
• 在模型里加 Dropout，对比是否减少过拟合；
• 引入数据增强，比如随机旋转一个小角度，测试对准确率的影响。

这些改动每一次只动一个变量，准确率的变化会被放大显现，最终你会形成对神经网络直觉性的理解。这个项目的意义不在于得到一个 98% 的准确率，而在于给你一个稳定的实验台，后续所有想法都可以在这个基础上快速验证。这也是我认为每个人第一个神经网络项目最应该达成的目标。