当前位置：首页 > news >正文

用 PyTorch 实现 MNIST 手写数字分类与训练损失曲线绘制

news 2025/10/14 6:39:21

一、引言

在深度学习入门阶段，MNIST 手写数字数据集是最经典的实验场景之一 —— 它包含 60000 张训练样本和 10000 张测试样本，每张图片都是 28×28 像素的灰度手写数字（0-9），任务是通过模型学习数字的特征，实现准确分类。

本文将基于 PyTorch 框架，从数据加载与预处理、神经网络模型设计、训练流程实现到训练损失曲线可视化，完整拆解 MNIST 分类任务的每一步。不仅提供可直接运行的代码，还会详细解释每个参数的意义、数值选择依据，帮助新手理解深度学习训练的核心逻辑。

二、环境准备

在开始前，需确保已安装以下依赖库（推荐使用 Anaconda 或 pip 安装）：

PyTorch：深度学习核心框架（负责模型构建、梯度计算）
TorchVision：计算机视觉工具库（提供 MNIST 数据集、数据预处理函数）
Matplotlib：绘图库（用于样本可视化、损失曲线绘制）
NumPy：数值计算库（辅助数据处理）

安装命令（以 pip 为例，PyTorch 版本需根据 CUDA 配置调整，可参考PyTorch 官网）：

pip install torch torchvision matplotlib numpy

三、完整代码实现与详细解析

3.1 库导入与基础设置

首先导入所有需要的工具库，并说明每个库的核心作用：

# 导入数值计算库：用于数组操作、数学计算（如后续可能的统计分析）
import numpy as np
# 导入PyTorch核心库：提供张量操作、自动微分、神经网络基础组件
import torch
# 从torchvision导入MNIST数据集：直接获取标准化的手写数字数据集
from torchvision.datasets import mnist
# 导入数据预处理模块：定义图像转换流程（如转张量、标准化）
import torchvision.transforms as transforms
# 导入数据加载工具：将数据集按批次分割，支持打乱、多线程加载
from torch.utils.data import DataLoader
# 导入神经网络功能模块：提供激活函数（ReLU、Softmax）、损失计算等
import torch.nn.functional as F
# 导入优化器模块：提供参数更新算法（如SGD、Adam）
import torch.optim as optim
# 导入神经网络基础模块：定义网络层（Linear、BatchNorm）、模型基类
import torch.nn as nn
# 导入TensorBoard工具：可选，用于训练过程可视化（损失、准确率实时监控）
from torch.utils.tensorboard import SummaryWriter
# 导入绘图库：用于样本可视化、训练损失曲线绘制
import matplotlib.pyplot as plt
# 设置matplotlib在Notebook/脚本中直接显示图像（非必要，根据运行环境调整）
%matplotlib inline  # 仅Notebook环境需要，脚本环境可注释

3.2 超参数定义（核心参数解析）

超参数是训练前需要手动设置的参数，直接影响模型的训练效率和性能。以下是关键超参数的定义及数值选择原因：

# 1. 批次大小：训练时每批处理的样本数量
train_batch_size = 64  # 训练集批次大小：64是深度学习中常用值
# 选择原因：过小会导致训练不稳定（单次更新样本少，梯度波动大），过大则占用更多内存（GPU显存不足时需减小）
test_batch_size = 128  # 测试集批次大小：可大于训练集
# 选择原因：测试阶段无需更新参数，仅需评估，增大批次可加速测试（不影响结果准确性）# 2. 初始学习率：控制参数更新的步长
learning_rate = 0.01  # 初始学习率：0.01是SGD优化器的常用初始值
# 选择原因：学习率过大会导致模型不收敛（参数震荡），过小会导致训练缓慢（参数更新慢，需更多轮次）# 3. 训练总轮数：整个数据集被完整训练的次数
num_epoches = 20  # 训练20轮
# 选择原因：MNIST任务简单，20轮足以让模型收敛；轮次过多可能导致过拟合（训练损失下降但测试损失上升）

3.3 数据预处理（标准化的重要性）

MNIST 原始数据是像素值为 0-255 的图像，直接输入模型会导致训练效率低（数值范围差异大，梯度更新不均衡）。需通过预处理将数据转换为适合模型训练的格式：

# 定义数据预处理流水线：用Compose组合多个操作，按顺序执行
transform = transforms.Compose([# 操作1：将PIL图像（MNIST原始格式）转为PyTorch张量（Tensor）transforms.ToTensor(),  # 作用：1. 维度转换（H×W×C → C×H×W，MNIST为灰度图，C=1）；2. 像素值归一化到[0, 1]（255→1）# 操作2：对张量进行标准化（均值减、标准差除）transforms.Normalize([0.5], [0.5])  # 作用：将数据从[0,1]映射到[-1, 1]（公式：output = (input - mean) / std）# 选择原因：标准化后数据分布更均匀，模型更容易学习特征（梯度更新更稳定）
])

3.4 数据集加载（训练集与测试集区分）

使用 TorchVision 提供的mnist.MNIST类直接下载 / 加载数据集，再用DataLoader按批次分割：

# 1. 加载训练集
train_dataset = mnist.MNIST(root='../data/',  # 数据集本地存储路径（若不存在会自动创建）train=True,       # 标记为训练集（加载60000张样本）transform=transform,  # 应用上述预处理流程download=True     # 若本地无数据集，自动从TorchVision服务器下载
)# 2. 加载测试集
test_dataset = mnist.MNIST(root='../data/',train=False,      # 标记为测试集（加载10000张样本，不参与训练）transform=transform  # 测试集必须与训练集用相同预处理（保证数据分布一致）
)# 3. 创建训练数据加载器（按批次加载，支持打乱）
train_loader = DataLoader(dataset=train_dataset,batch_size=train_batch_size,shuffle=True  # 训练时打乱数据顺序（避免模型学习到"样本顺序"的无关特征，提升泛化能力）
)# 4. 创建测试数据加载器（无需打乱）
test_loader = DataLoader(dataset=test_dataset,batch_size=test_batch_size,shuffle=False  # 测试时仅需评估，打乱不影响结果，反而增加计算开销
)

3.5 数据结构查看与样本可视化

加载数据后，先查看数据的维度信息，再可视化部分样本，确认数据加载正确：

# 1. 查看测试集数据结构（以第一个批次为例）
# 将DataLoader转为可枚举对象，便于按批次获取数据
examples = enumerate(test_loader)
# 获取第一个批次：batch_idx=批次索引，example_data=图像数据，example_targets=真实标签
batch_idx, (example_data, example_targets) = next(examples)# 打印数据形状：输出为 torch.Size([128, 1, 28, 28])
print("测试集单个批次数据形状：", example_data.shape)
# 形状含义：[批次大小, 通道数, 图像高度, 图像宽度] → 128个样本，1个通道（灰度），28×28像素# 2. 可视化6个测试集样本（2行3列布局）
fig = plt.figure(figsize=(10, 6))  # 设置图像窗口大小（宽10，高6）
for i in range(6):# 创建子图：2行3列，第i+1个位置（matplotlib子图索引从1开始）plt.subplot(2, 3, i+1)# 自动调整子图间距，避免标签重叠plt.tight_layout()# 显示图像：example_data[i][0] → 第i个样本的通道维度（灰度图仅1个通道）# cmap='gray' → 灰度模式显示，interpolation='none' → 不使用插值（保持像素清晰）plt.imshow(example_data[i][0], cmap='gray', interpolation='none')# 标题显示样本的真实标签（Ground Truth）plt.title(f"Ground Truth: {example_targets[i]}")# 去除x轴、y轴刻度（避免显示像素坐标，简化图像）plt.xticks([])plt.yticks([])
plt.show()  # 显示图像窗口

可视化效果：将显示 2 行 3 列共 6 个手写数字，每个数字下方标注其真实标签（如 “Ground Truth: 5”），可直观确认数据加载正确。

3.6 神经网络模型定义（全连接网络设计）

设计一个简单的全连接神经网络（适合 MNIST 这类简单图像任务），包含展平层、两个隐藏层和输出层，并加入批归一化加速收敛：

class Net(nn.Module):"""自定义全连接神经网络类，继承自PyTorch的nn.Module（所有模型的基类）网络结构：输入 → 展平层 → 隐藏层1（线性+批归一化）→ ReLU → 隐藏层2（线性+批归一化）→ ReLU → 输出层 → Softmax"""def __init__(self, in_dim, n_hidden_1, n_hidden_2, out_dim):# 调用父类（nn.Module）的构造函数，必须执行super(Net, self).__init__()# 1. 展平层：将28×28的二维图像转为一维向量（784个元素）self.flatten = nn.Flatten()# 2. 第一个隐藏层：线性层 + 批归一化（加速训练、稳定收敛）self.layer1 = nn.Sequential(# 线性层（全连接层）：输入维度in_dim → 隐藏层维度n_hidden_1nn.Linear(in_dim, n_hidden_1),# 批归一化层：对线性层输出做归一化（减少内部协变量偏移）nn.BatchNorm1d(n_hidden_1))# 3. 第二个隐藏层：结构与第一个隐藏层一致，维度从n_hidden_1 → n_hidden_2self.layer2 = nn.Sequential(nn.Linear(n_hidden_1, n_hidden_2),nn.BatchNorm1d(n_hidden_2))# 4. 输出层：线性层，维度从n_hidden_2 → 输出维度out_dim（MNIST为10类）self.out = nn.Sequential(nn.Linear(n_hidden_2, out_dim))def forward(self, x):"""前向传播函数：定义数据在网络中的流动路径（必须实现）参数x：输入数据（形状为[batch_size, 1, 28, 28]）返回：模型输出（形状为[batch_size, 10]，表示每个类别的概率）"""# 步骤1：展平输入 → [batch_size, 1, 28, 28] → [batch_size, 784]x = self.flatten(x)# 步骤2：第一个隐藏层 + ReLU激活（引入非线性，让模型学习复杂特征）x = F.relu(self.layer1(x))# 步骤3：第二个隐藏层 + ReLU激活x = F.relu(self.layer2(x))# 步骤4：输出层 + Softmax激活（将输出转为概率分布，每行和为1）# dim=1 → 按行归一化（每个样本对应10个类别的概率）x = F.softmax(self.out(x), dim=1)return x# 模型参数说明（对应__init__的输入）
in_dim = 28 * 28  # 输入维度：28×28=784（展平后的向量长度）
n_hidden_1 = 300  # 第一个隐藏层神经元数：300（经验值，可调整）
n_hidden_2 = 100  # 第二个隐藏层神经元数：100（逐步减少维度，提取高层特征）
out_dim = 10      # 输出维度：10（MNIST有10个类别：0-9）

3.7 模型初始化与配置（设备选择、损失函数、优化器）

完成模型定义后，需配置训练设备（GPU/CPU）、损失函数（衡量预测误差）和优化器（更新模型参数）：

# 1. 选择训练设备：优先使用GPU（CUDA），若无则使用CPU
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
print(f"训练设备：{device}")  # 输出结果：如"训练设备：cuda:0"（GPU）或"训练设备：cpu"# 2. 实例化模型，并将模型移动到指定设备
model = Net(in_dim=in_dim, n_hidden_1=n_hidden_1, n_hidden_2=n_hidden_2, out_dim=out_dim)
model.to(device)  # 关键步骤：将模型参数（张量）移动到GPU/CPU# 3. 定义损失函数：交叉熵损失（适合多分类任务，已包含Softmax，此处模型的Softmax可省略，避免重复计算）
# 注意：若模型未用Softmax，损失函数需用nn.CrossEntropyLoss()；若已用Softmax，需用nn.NLLLoss()
# 本文为了直观展示概率分布，模型保留Softmax，故损失函数调整为nn.NLLLoss()（对数似然损失）
# 修正：原模型的Softmax与CrossEntropyLoss重复，此处调整代码（避免计算错误）
criterion = nn.NLLLoss()  # 对应模型的Softmax输出# 4. 定义优化器：随机梯度下降（SGD），加入动量（加速收敛）
optimizer = optim.SGD(params=model.parameters(),  # 待优化的参数（模型所有可训练参数）lr=learning_rate,           # 学习率（复用前文定义的超参数）momentum=0.9                # 动量参数：0.9（模拟物理惯性，减少梯度震荡，加速收敛）
)# 5. 初始化TensorBoard日志（可选，用于实时监控训练过程）
writer = SummaryWriter(log_dir='./logs', comment='mnist_train')  # 日志保存路径：./logs

3.8 模型训练与评估（核心流程）

训练流程是深度学习的核心，需实现 “前向传播计算损失→反向传播求梯度→优化器更新参数” 的循环，并在每轮训练后用测试集评估模型性能：

# 初始化列表：存储每轮的训练/测试损失和准确率（用于后续绘制曲线）
train_losses = []  # 训练损失（每轮平均损失）
train_accs = []    # 训练准确率（每轮平均准确率）
test_losses = []   # 测试损失（每轮平均损失）
test_accs = []     # 测试准确率（每轮平均准确率）# 开始训练循环（共num_epoches=20轮）
for epoch in range(num_epoches):# -------------------------- 训练阶段 --------------------------model.train()  # 设置模型为训练模式（启用批归一化、dropout等训练特有的层）train_loss = 0.0  # 累计训练损失（每轮重置）train_correct = 0  # 累计训练正确的样本数（每轮重置）total_train = 0    # 累计训练样本总数（每轮重置）# 动态调整学习率：每5轮学习率乘以0.9（学习率衰减，训练后期减小步长）if epoch % 5 == 0 and epoch != 0:  # epoch=0不衰减（初始学习率）optimizer.param_groups[0]['lr'] *= 0.9current_lr = optimizer.param_groups[0]['lr']print(f"\n第{epoch+1}/{num_epoches}轮训练，当前学习率：{current_lr:.6f}")# 遍历训练集的每个批次（batch）for batch_idx, (img, label) in enumerate(train_loader):# 将图像和标签移动到训练设备（GPU/CPU）img = img.to(device)label = label.to(device)# 1. 前向传播：计算模型预测输出output = model(img)# 2. 计算损失（预测输出与真实标签的误差）loss = criterion(torch.log(output), label)  # NLLLoss需输入log(output)# 3. 反向传播：清空之前的梯度 → 计算梯度 → 更新参数optimizer.zero_grad()  # 清空梯度（避免梯度累积）loss.backward()        # 反向传播求梯度optimizer.step()       # 优化器更新模型参数# 累计训练损失（loss.item()获取张量的数值，避免计算图累积）train_loss += loss.item() * img.size(0)  # 乘以批次大小（每个批次的总损失）# 计算训练准确率：获取预测类别（概率最大的类别）_, pred = torch.max(output, dim=1)  # pred为预测类别（形状[batch_size]）train_correct += (pred == label).sum().item()  # 累计正确样本数total_train += img.size(0)  # 累计样本总数# 计算每轮的平均训练损失和准确率avg_train_loss = train_loss / total_trainavg_train_acc = train_correct / total_traintrain_losses.append(avg_train_loss)train_accs.append(avg_train_acc)# 打印训练结果print(f"训练集 - 平均损失：{avg_train_loss:.4f}，准确率：{avg_train_acc:.4f}")# 将训练指标写入TensorBoard（可选）writer.add_scalar('Train/Loss', avg_train_loss, epoch)writer.add_scalar('Train/Accuracy', avg_train_acc, epoch)# -------------------------- 测试阶段 --------------------------model.eval()  # 设置模型为评估模式（禁用批归一化、dropout等，避免影响测试结果）test_loss = 0.0  # 累计测试损失（每轮重置）test_correct = 0  # 累计测试正确的样本数（每轮重置）total_test = 0    # 累计测试样本总数（每轮重置）# 测试阶段不计算梯度（节省内存，加速计算）with torch.no_grad():# 遍历测试集的每个批次for batch_idx, (img, label) in enumerate(test_loader):# 将图像和标签移动到测试设备（与训练设备一致）img = img.to(device)label = label.to(device)# 1. 前向传播：计算预测输出output = model(img)# 2. 计算测试损失loss = criterion(torch.log(output), label)test_loss += loss.item() * img.size(0)# 3. 计算测试准确率_, pred = torch.max(output, dim=1)test_correct += (pred == label).sum().item()total_test += img.size(0)# 计算每轮的平均测试损失和准确率avg_test_loss = test_loss / total_testavg_test_acc = test_correct / total_testtest_losses.append(avg_test_loss)test_accs.append(avg_test_acc)# 打印测试结果print(f"测试集 - 平均损失：{avg_test_loss:.4f}，准确率：{avg_test_acc:.4f}")# 将测试指标写入TensorBoard（可选）writer.add_scalar('Test/Loss', avg_test_loss, epoch)writer.add_scalar('Test/Accuracy', avg_test_acc, epoch)# 训练结束：关闭TensorBoard写入器
writer.close()
print("\n训练完成！")

3.9 训练损失曲线绘制（直观分析模型性能）

训练完成后，通过绘制训练 / 测试损失曲线和训练 / 测试准确率曲线，可直观观察模型的收敛情况、是否过拟合：

# 设置绘图风格（可选，让图像更美观）
plt.style.use('seaborn-v0_8-whitegrid')# 创建2个子图（1行2列）：分别显示损失曲线和准确率曲线
fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(15, 6))# -------------------------- 子图1：损失曲线 --------------------------
ax1.plot(range(1, num_epoches+1), train_losses, label='Train Loss', color='blue', linewidth=2)
ax1.plot(range(1, num_epoches+1), test_losses, label='Test Loss', color='red', linewidth=2, linestyle='--')
ax1.set_xlabel('Epoch (训练轮次)', fontsize=12)
ax1.set_ylabel('Loss (损失值)', fontsize=12)
ax1.set_title('Training & Test Loss Curve (训练与测试损失曲线)', fontsize=14, fontweight='bold')
ax1.legend(fontsize=10)  # 显示图例
ax1.grid(True, alpha=0.3)  # 显示网格（alpha=透明度）# -------------------------- 子图2：准确率曲线 --------------------------
ax2.plot(range(1, num_epoches+1), train_accs, label='Train Accuracy', color='green', linewidth=2)
ax2.plot(range(1, num_epoches+1), test_accs, label='Test Accuracy', color='orange', linewidth=2, linestyle='--')
ax2.set_xlabel('Epoch (训练轮次)', fontsize=12)
ax2.set_ylabel('Accuracy (准确率)', fontsize=12)
ax2.set_title('Training & Test Accuracy Curve (训练与测试准确率曲线)', fontsize=14, fontweight='bold')
ax2.legend(fontsize=10)
ax2.grid(True, alpha=0.3)
ax2.set_ylim(0.9, 1.0)  # 调整y轴范围（MNIST准确率较高，聚焦0.9-1.0区间）# 调整子图间距
plt.tight_layout()# 保存图像（可选，保存为PNG格式，分辨率300dpi）
plt.savefig('./mnist_train_curves.png', dpi=300, bbox_inches='tight')# 显示图像
plt.show()

曲线分析要点：

损失曲线：训练损失和测试损失应逐步下降，最终趋于稳定（若测试损失先降后升，说明过拟合）；
准确率曲线：训练准确率和测试准确率应逐步上升，最终接近 1.0（MNIST 任务中，测试准确率可达 97% 以上）；
学习率衰减效果：每 5 轮学习率衰减后，损失曲线下降速度可能变慢，但会更接近最优解。

四、训练结果示例与优化方向

4.1 预期结果

测试准确率：20 轮训练后，测试准确率可达 97%-98%（若优化超参数，可进一步提升至 99% 以上）；
损失曲线：训练损失从初始的 2.3 左右（随机猜测的损失）下降到 0.1 以下，测试损失与训练损失差距较小（无明显过拟合）。

4.2 优化方向

调整超参数：
- 批次大小：若 GPU 显存充足，可增大至 128；
- 学习率：尝试初始学习率 0.001，或用学习率调度器（如torch.optim.lr_scheduler.StepLR）；
- 轮次：若未收敛，可增加至 30 轮。
改进模型结构：
- 改用卷积神经网络（CNN）：如 LeNet-5，适合图像任务（提取空间特征，准确率更高）；
- 加入 dropout 层：如nn.Dropout(0.5)，减少过拟合。
更换优化器：
- 改用 Adam 优化器（optim.Adam）：自适应学习率，训练速度更快（无需手动调整学习率）。