PyTorch实战：手把手教你用简单神经网络识别MNIST手写数字（附完整代码）

在前面的文章中，我们已经学会了用PyTorch加载MNIST数据集并进行可视化。今天，我们将更进一步——用神经网络模型训练一个能自动识别手写数字的分类器。本文会尽量简化理论，聚焦代码实操，帮你快速上手深度学习模型训练全流程！

---

一、目标明确：我们要做什么？

我们的目标是让计算机“看懂”手写数字图片（比如分辨“3”和“5”）。具体步骤如下：

1. 准备数据：用PyTorch加载MNIST训练集和测试集。
2. 搭建模型：定义一个简单的神经网络（类似“信息加工流水线”）。
3. 训练模型：用训练数据“教”模型识别数字规律。
4. 测试模型：用测试数据验证模型的“实战能力”。

---

二、环境准备：设备与数据加载

1. 设备选择：优先用GPU加速

深度学习需要大量计算，GPU（显卡）比CPU快很多。PyTorch中可以用一行代码指定设备：

device = 'cuda' if torch.cuda.is_available() else 'cpu'  # 自动检测是否有GPU
print(f"当前使用设备：{device}")  # 输出：cuda（如果有GPU）或cpu

• 如果你有NVIDIA显卡（常见游戏本/工作站），cuda会启用GPU加速；没有的话自动用CPU（速度慢但也能跑）。

2. 数据加载：用DataLoader批量管理数据

之前我们已经下载了MNIST数据集，现在需要用DataLoader按“批次”加载数据（每次训练64张图）：

from torch.utils.data import DataLoader

# 训练集：打乱顺序（shuffle=True），每次取64张
train_dataloader = DataLoader(training_data, batch_size=64, shuffle=True)
# 测试集：不打乱顺序（shuffle=False），每次取64张
test_dataloader = DataLoader(test_data, batch_size=64, shuffle=False)

• batch_size=64：每次训练64张图（称为一个“批次”），太小会浪费计算资源，太大会导致内存不足。
• shuffle=True：训练时打乱数据顺序，避免模型“记住”数据的排列顺序（比如前64张都是“1”）。

---

三、搭建神经网络：定义“信息加工流水线”

神经网络的核心是**层（Layer）**的堆叠，每一层负责将输入数据转换为更有用的特征。我们搭建一个简单的全连接神经网络（适合新手理解）：

import torch.nn as nn

class HandwrittenDigitClassifier(nn.Module):  # 继承PyTorch的神经网络基类
    def __init__(self):
        super().__init__()  # 调用父类初始化方法（必须写）
        
        # 层1：将28x28的图片展平为1维向量（784个元素）
        self.flatten = nn.Flatten()
        # 层2：全连接层1（输入784，输出128）—— 学习基础特征（如线条、拐点）
        self.hidden1 = nn.Linear(28*28, 128)
        # 层3：全连接层2（输入128，输出256）—— 学习更复杂特征（如数字整体结构）
        self.hidden2 = nn.Linear(128, 256)
        # 层4：输出层（输入256，输出10）—— 对应0-9十个数字的分类结果
        self.out = nn.Linear(256, 10)

    def forward(self, x):
        # 前向传播：定义数据如何通过各层
        x = self.flatten(x)  # 输入形状：[64,1,28,28] → 展平为[64,784]（64是批次大小）
        x = self.hidden1(x)  # 全连接层1：[64,784] → [64,128]
        x = torch.sigmoid(x)  # 激活函数：给数据“加非线性”（避免模型只会学直线）
        x = self.hidden2(x)  # 全连接层2：[64,128] → [64,256]
        x = torch.sigmoid(x)  # 再次激活
        x = self.out(x)       # 输出层：[64,256] → [64,10]（10个数字的概率）
        return x

代码通俗解释：

• nn.Flatten()：把图片的“二维像素矩阵”（28x28）和“通道”（1）展成一维向量（784个数字），因为全连接层只能处理一维数据。
• nn.Linear(…)：全连接层，每个神经元与前一层所有神经元相连。例如hidden1层有128个神经元，每个神经元会“接收”前一层784个神经元的信息，然后输出128个新信息。
• torch.sigmoid(x)：激活函数，把输出值压缩到0-1之间。它的作用是让模型能学习“非线性关系”（比如数字“8”的上下两个圈不是直线）。

---

四、训练准备：损失函数与优化器

1. 损失函数：衡量预测值与真实值的差距

我们用交叉熵损失（CrossEntropyLoss），它专门用于多分类任务（比如0-9十个类别）。简单说，它的作用是告诉模型：“你这次预测错了多少，需要调整参数来减少错误”。

loss_fn = nn.CrossEntropyLoss()  # 定义损失函数

2. 优化器：根据损失调整模型参数

优化器的作用是“根据损失函数的反馈，调整模型中的参数（比如全连接层的权重）”。我们用最基础的随机梯度下降（SGD）：

optimizer = torch.optim.SGD(model.parameters(), lr=0.01)  # 学习率lr=0.01

• model.parameters()：告诉优化器要调整哪些参数（模型中所有可学习的权重和偏置）。
• lr=0.01：学习率，控制参数调整的步长。如果lr太大（比如1），模型可能“跳过”正确的参数；如果lr太小（比如0.0001），训练会很慢。

---

五、训练循环：让模型“学习”规律

训练过程就像教小孩认字：先给例子（数据），计算错误（损失），然后调整（优化）。我们把这个过程写成一个train函数：

def train(dataloader, model, loss_fn, optimizer):
    model.train()  # 开启训练模式（某些层会启用随机行为，如Dropout）
    total_loss = 0  # 记录总损失
    correct = 0     # 记录正确预测的数量
    total = len(dataloader.dataset)  # 总样本数（6万）

    for batch, (X, y) in enumerate(dataloader):  # 遍历每个批次（共6万/64≈938批）
        X, y = X.to(device), y.to(device)  # 把数据和标签搬到GPU（如果可用）

        # 1. 前向传播：模型预测
        pred = model(X)  # 输入X（64张图），输出pred（64个预测结果，每个是10维向量）

        # 2. 计算损失：预测值与真实值的差距
        loss = loss_fn(pred, y)  # 损失是一个标量（单个数字）
        total_loss += loss.item()  # 累加每批次的损失

        # 3. 反向传播：计算参数的梯度（误差的“方向”）
        optimizer.zero_grad()  # 清空之前的梯度（避免累积）
        loss.backward()        # 反向传播计算当前梯度

        # 4. 参数更新：优化器根据梯度调整参数
        optimizer.step()       # 按梯度方向调整参数（让损失变小）

        # 统计正确预测数
        correct += (pred.argmax(1) == y).type(torch.float).sum().item()

        # 每100个批次打印一次训练进度
        if batch % 100 == 0:
            current_batch = batch + 1  # 当前是第几个批次（从1开始）
            print(f"批次 {current_batch}/938 | 平均损失: {total_loss/current_batch:.4f} | 正确率: {100*correct/(current_batch*64):.2f}%")

关键步骤解释：

• model.train()：告诉模型“现在是训练时间”，一些特殊层（如Dropout）会随机“关闭”部分神经元，避免模型“死记硬背”训练数据。
• X.to(device), y.to(device)：把数据和标签从CPU搬到GPU（如果可用），利用GPU的并行计算加速训练。
• optimizer.zero_grad()：清空之前计算的梯度（PyTorch会累积梯度，多次反向传播会累加，所以每次更新前必须清零）。
• loss.backward()：反向传播计算每个参数的梯度（即损失对参数的偏导数），梯度越大，说明参数离最优值越远。
• optimizer.step()：优化器根据梯度调整参数（比如SGD的公式：参数 = 参数 - lr × 梯度），让损失变小。
• pred.argmax(1)：找到预测结果中概率最大的那个数字（比如pred是[0.1, 0.3, …, 0.6]，argmax会返回2，表示预测为2）。

---

六、测试评估：验证模型的“实战能力”

训练完成后，我们需要用测试集（模型没见过的数据）评估它的泛化能力（对新数据的识别能力）。我们写一个test函数计算准确率：

def test(dataloader, model, loss_fn):
    model.eval()  # 开启测试模式（关闭Dropout等随机层）
    total_loss = 0
    correct = 0
    total = len(dataloader.dataset)  # 总样本数（1万）

    with torch.no_grad():  # 关闭梯度计算（测试不需要反向传播，省内存）
        for X, y in dataloader:
            X, y = X.to(device), y.to(device)
            pred = model(X)  # 模型预测
            total_loss += loss_fn(pred, y).item()  # 累加损失
            correct += (pred.argmax(1) == y).type(torch.float).sum().item()  # 统计正确数

    # 计算平均损失和准确率
    avg_loss = total_loss / len(dataloader)
    accuracy = correct / total
    print(f"测试结果 | 平均损失: {avg_loss:.4f} | 准确率: {accuracy*100:.2f}%\n")
    return accuracy

关键细节：

• model.eval()：告诉模型“现在是测试时间”，关闭Dropout等随机层（测试时不需要随机，要保持稳定）。
• with torch.no_grad()：上下文管理器，关闭梯度计算。测试时不需要反向传播，关闭梯度可以节省内存。

---

七、完整代码与运行结果

现在，我们把所有代码整合起来，运行看看效果：

import torch
from torch import nn
from torch.utils.data import DataLoader
from torchvision import datasets
from torchvision.transforms import ToTensor
from matplotlib import pyplot as plt

# ---------------------- 环境准备 ----------------------
device = 'cuda' if torch.cuda.is_available() else 'cpu'
print(f"当前使用设备：{device}")

# 加载数据（之前已经下载过，download=False）
training_data = datasets.MNIST(root="data", train=True, download=False, transform=ToTensor())
test_data = datasets.MNIST(root="data", train=False, download=False, transform=ToTensor())

# 数据加载器
train_dataloader = DataLoader(training_data, batch_size=64, shuffle=True)
test_dataloader = DataLoader(test_data, batch_size=64, shuffle=False)

# ---------------------- 模型定义 ----------------------
class HandwrittenDigitClassifier(nn.Module):
    def __init__(self):
        super().__init__()
        self.flatten = nn.Flatten()
        self.hidden1 = nn.Linear(28*28, 128)
        self.hidden2 = nn.Linear(128, 256)
        self.out = nn.Linear(256, 10)

    def forward(self, x):
        x = self.flatten(x)
        x = self.hidden1(x)
        x = torch.sigmoid(x)
        x = self.hidden2(x)
        x = torch.sigmoid(x)
        x = self.out(x)
        return x

# ---------------------- 训练准备 ----------------------
model = HandwrittenDigitClassifier().to(device)  # 模型移到GPU
loss_fn = nn.CrossEntropyLoss()  # 损失函数
optimizer = torch.optim.SGD(model.parameters(), lr=0.01)  # 优化器（lr=0.01更合理）

# ---------------------- 训练与测试 ----------------------
def train(dataloader, model, loss_fn, optimizer):
    model.train()
    total_loss = 0
    correct = 0
    total = len(dataloader.dataset)

    for batch, (X, y) in enumerate(dataloader):
        X, y = X.to(device), y.to(device)
        pred = model(X)
        loss = loss_fn(pred, y)
        total_loss += loss.item()
        correct += (pred.argmax(1) == y).type(torch.float).sum().item()

        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        if batch % 100 == 0:
            current_batch = batch + 1
            print(f"批次 {current_batch}/938 | 平均损失: {total_loss/current_batch:.4f} | 正确率: {100*correct/(current_batch*64):.2f}%")

    avg_loss = total_loss / len(dataloader)
    accuracy = correct / total
    print(f"训练完成 | 平均损失: {avg_loss:.4f} | 整体正确率: {accuracy*100:.2f}%\n")

def test(dataloader, model, loss_fn):
    model.eval()
    total_loss = 0
    correct = 0
    total = len(dataloader.dataset)

    with torch.no_grad():
        for X, y in dataloader:
            X, y = X.to(device), y.to(device)
            pred = model(X)
            total_loss += loss_fn(pred, y).item()
            correct += (pred.argmax(1) == y).type(torch.float).sum().item()

    avg_loss = total_loss / len(dataloader)
    accuracy = correct / total
    print(f"测试结果 | 平均损失: {avg_loss:.4f} | 准确率: {accuracy*100:.2f}%\n")
    return accuracy

# 训练10轮（Epoch）
num_epochs = 10
for epoch in range(num_epochs):
    print(f"===== 第 {epoch+1}/{num_epochs} 轮训练 =====")
    train(train_dataloader, model, loss_fn, optimizer)
    test(test_dataloader, model, loss_fn)

# 最终测试（可选）
final_acc = test(test_dataloader, model, loss_fn)
print(f"最终测试准确率：{final_acc*100:.2f}%")

运行结果示例（假设使用GPU）：

当前使用设备：cuda
===== 第 1/10 轮训练 =====
批次 1/938 | 平均损失: 2.3026 | 正确率: 10.94%
...
批次 938/938 | 平均损失: 0.3452 | 正确率: 89.06%
训练完成 | 平均损失: 0.6823 | 整体正确率: 85.23%
测试结果 | 平均损失: 0.3215 | 准确率: 89.12%

===== 第 2/10 轮训练 =====
...
===== 第 10/10 轮训练 =====
测试结果 | 平均损失: 0.2104 | 准确率: 93.45%
最终测试准确率：93.45%

---

八、常见问题与改进建议

1. 学习率太大怎么办？

如果训练时损失突然变大（比如从2.3跳到5.0），说明学习率（lr）太大，模型“步子迈太大”导致无法收敛。可以尝试调小lr（比如从0.01降到0.001）。

2. 模型准确率上不去？

• 调整模型复杂度：增加隐藏层（比如加一个hidden3层）或增大隐藏层神经元数量（比如hidden1=256）。
• 更换激活函数：把sigmoid换成ReLU（nn.ReLU()），它能缓解梯度消失问题，加速训练。
• 增加训练轮次：训练10轮可能不够，试试20轮（num_epochs=20）。

3. GPU没生效？

检查device是否正确（print(device)应输出cuda），确保PyTorch已安装CUDA支持（可通过torch.cuda.is_available()验证）。

---

九、总结

通过本文，我们完成了从数据加载到模型训练、测试的全流程，实现了MNIST手写数字分类。关键步骤包括：

• 设备配置：优先用GPU加速训练；
• 模型搭建：用nn.Module定义神经网络结构；
• 训练循环：前向传播计算损失，反向传播更新参数；
• 测试评估：用测试集验证模型泛化能力。

深度学习的核心是“数据+模型+训练”，多动手调整参数（比如lr、隐藏层大小）、观察结果变化，你会逐渐掌握其中的规律。下次可以尝试：

• 用卷积神经网络（CNN）提升准确率（CNN更适合图像任务）；
• 添加Dropout层防止过拟合；
• 调整batch_size（比如从64改为128）观察训练稳定性。