机器学习详解(5):MLP代码详解之MNIST手写数字识别

在上一篇文章 机器学习详解(4):多层感知机MLP之理论学习中，我们学习了MLP的理论。在深度学习中，MNIST手写数字数据集被誉为“深度学习的Hello World”，在图像分类问题中极具代表性。本文将基于一段简单的Python代码，一起来学习如何使用多层感知器（MLP）来完成手写数字的分类任务。

1 MNIST数据集

MNIST（Modified National Institute of Standards and Technology）是一个经典的手写数字数据集，包含以下特点：

• 数据内容：共70,000张28x28像素的灰度图像，其中包括60,000张用于训练的数据和10,000张用于测试的数据。
• 标签分类：每张图片对应一个从0到9的数字标签，共10个类别。
• 任务目标：构建一个分类模型，使其能够根据输入图像准确预测数字的类别。

下面是40张来自于MNIST数据集的图片：

MNIST数据集的意义在于其广泛的使用和相对简单的特性。作为许多深度学习算法的基准测试集，它让研究者能够快速验证模型的性能。此外，由于数据规模较小，模型可以快速训练和测试，非常适合入门学习和实验验证。

2 代码详解

2.1 导入库和GPU

1.导入需要使用的库

import torch
import torch.nn as nn
from torch.utils.data import Dataset, DataLoader
from torch.optim import Adam# Visualization tools
import torchvision
import torchvision.transforms.v2 as transforms
import torchvision.transforms.functional as F
import matplotlib.pyplot as plt

• torch：PyTorch的核心库，用于构建和训练深度学习模型，支持张量操作、自动微分等功能。
• torch.nn：PyTorch的神经网络模块，提供了常用的神经网络层（如全连接层、卷积层）和相关功能（如激活函数、损失函数）。
• torch.utils.data.Dataset：数据加载工具，用于定义自定义数据集类，实现数据的加载与预处理。
• torch.utils.data.DataLoader：数据加载器，结合Dataset，用于按批次加载数据并支持多线程。
• torch.optim.Adam：PyTorch优化器模块，Adam是常用的优化算法之一，用于调整模型的参数以最小化损失函数。
• torchvision：PyTorch的计算机视觉工具包，包含常用的数据集、模型和图像处理工具。
• torchvision.transforms.v2：图像变换模块（新版），提供用于图像预处理的功能，如归一化、裁剪、旋转等。
• torchvision.transforms.functional：功能性图像变换模块，提供粒度更细的操作，如手动指定每个步骤的参数。
• matplotlib.pyplot：Python的可视化库，用于绘制图表或可视化模型的训练过程、数据分布等。

2.GPU设置

在PyTorch中，我们可以通过将设备设置为cuda来在GPU上运行操作。函数torch.cuda.is_available()会验证PyTorch是否能识别GPU。

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
torch.cuda.is_available()

2.2 MNIST数据集处理

2.2.1 下载和导入

我们需要为MNIST数据集准备4个数据片段：

• x_train: 用于训练神经网络的图像数据。
• y_train: 与x_train图像对应的正确标签，用于评估模型在训练过程中的预测结果。
• x_valid: 为验证模型性能而预留的图像数据，模型训练完成后使用。
• y_valid: 与x_valid图像对应的正确标签，用于评估模型在训练完成后的预测结果。

MNIST 数据集可以通过 PyTorch 的 TorchVision 库直接下载并加载，这大大简化了数据管理的流程。

train_set = torchvision.datasets.MNIST("./data/", train=True, download=True)
valid_set = torchvision.datasets.MNIST("./data/", train=False, download=True)

• 使用 download=True 参数时，如果指定路径下没有数据集，TorchVision 会下载数据并存储在 ./data/ 目录。

• 数据集被分为训练集（train=True）和验证集（train=False）。

可以发现 TorchVision 将其中 60,000 张图像划分为训练集，10,000 张图像划分为验证集（训练后用于验证）。

train_set
输出:
Dataset MNISTNumber of datapoints: 60000Root location: ./data/Split: Trainvalid_set
输出:
Dataset MNISTNumber of datapoints: 10000Root location: ./data/Split: Test

接着我们输出一下训练集的内容：

x_0, y_0 = train_set[0]

其中y_0为其对应的数字结果5，x_0为手写数字的图片：

2.2.2 张量(Tensors)

GPU 在张量处理方面高效，因为它具有大量并行计算核心，可以同时执行成千上万个简单数学操作，这非常适合处理多维数组（张量）的计算。再加上 GPU 专为矩阵运算优化的硬件架构，它能快速完成神经网络中常见的张量操作，例如矩阵乘法和加法。

接下来，我们将把图像转换为张量，以便后续用神经网络进行处理。TorchVision 提供了一个非常实用的工具类 ToTensor，可将 PIL 图像转换为张量格式。

trans = transforms.Compose([transforms.ToTensor()])
x_0_tensor = trans(x_0)x_0_tensor.dtype
输出:torch.float32

• Compose 需要接收一个列表，列表中包含一组按顺序执行的转换操作，这里表示只有一个列表

PIL 图像的像素值范围为整数 [0, 255]，但 ToTensor 类会将其转换为浮点数范围 [0.0, 1.0]。

x_0_tensor.min()
输出:
tensor(0.)x_0_tensor.max()
输出:
tensor(1.)

我们还可以查看每个维度的大小。PyTorch有三个维度(颜色通道，高度和宽度)$C\times H\times W$。由于这些图像是黑白的，因此只有 1 个颜色通道。图像是正方形，高度和宽度均为 28 像素。

x_0_tensor.size()
输出:
torch.Size([1, 28, 28])

默认情况下，张量是在 CPU 上处理的。

x_0_tensor.device
输出:
device(type='cpu')

如果需要将其移动到 GPU，可以使用 .cuda 方法。

x_0_gpu = x_0_tensor.cuda()
x_0_gpu.device
输出:
device(type='cuda', index=0)

需要注意的是，如果 PyTorch 未识别到 GPU，.cuda 方法将会失败。为了确保代码能在不同设备上灵活运行，我们可以使用 .to(device) 方法将张量移动到系统检测到的设备（如 GPU 或 CPU）。随后，通过 .device 属性检查张量当前所在的设备，确保其已正确迁移。

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
x_0_tensor.to(device).device
输出:
device(type='cuda', index=0)

有时，直接解读大量数值可能会很困难。幸运的是，TorchVision 提供了 to_pil_image 函数，可以将 $C\times H\times W$ 格式的张量转换回 PIL 图像。

image = F.to_pil_image(x_0_tensor)
plt.imshow(image, cmap='gray')

2.2.3 准备训练数据

1. 转换操作（Transforms）
转换（Transforms）是 torchvision 提供的一组函数，用于对数据集进行变换操作。例如，将图像转换为张量。
使用 Compose 组合转换函数可以将多个转换操作组合在一起：

trans = transforms.Compose([transforms.ToTensor()])

这里定义了一个简单的转换，将图像从 PIL 格式转换为张量。转换可以直接应用于单个数据点，也可以设置为数据集的 transform 属性，对整个数据集进行批量转换：

train_set.transform = trans
valid_set.transform = trans

2.数据加载器（DataLoaders）

数据加载器（DataLoader）定义了如何从数据集中取出数据用于训练模型。它可以按批次（batch）加载数据，方便高效地训练模型。

批量训练（Batch Training）：按批次训练模型不仅节省计算资源，还能提高模型训练效率。

• 批量大小(Batch Size)：通常设置为 32 或 64，批量大小过大会耗尽内存，过小则可能影响模型学习效率。

训练数据（train_loader）：需要随机打乱数据（shuffle=True）以避免模型过拟合到数据顺序。

验证数据（valid_loader）：无需打乱数据，但仍按批次加载以节省内存。

batch_size = 32train_loader = DataLoader(train_set, batch_size=batch_size, shuffle=True)
valid_loader = DataLoader(valid_set, batch_size=batch_size)

2.3 创建模型

我们将构建一个简单的 MLP 模型，每一层对接收到的数据进行数学运算后传递给下一层，包含以下四部分：

1. Flatten 层：将 n 维数据（如图像数据）转换为一维向量，作为 MLP 的输入。
2. 输入层：MLP 的第一层神经元，用于接收展开后的数据。
3. 隐藏层：MLP 的中间层，包含若干神经元，用于提取特征和表示。
4. 输出层：MLP 的最后一层神经元，生成模型的最终预测结果。

2.3.1 图像展开

输入数据通常是 3 维张量$C\times H\times W$，如灰度图像为 1×28×28。为了输入 MLP，需要将其转为 1 维向量（例如 1x784）。来看一个例子：

test_matrix = torch.tensor([[1, 2, 3],[4, 5, 6],[7, 8, 9]]
)
print(test_matrix)print(n.Flatten()(test_matrix))输出:
tensor([[1, 2, 3],[4, 5, 6],[7, 8, 9]])tensor([[1, 2, 3],[4, 5, 6],[7, 8, 9]])

注意：此时 Flatten 并未生效，因为神经网络期望输入的是一个批次数据（batch）。目前，Flatten 层需要输入一个 3 行向量作为单独的样本，而不是一个 2D 矩阵。

批量处理（Batching the Data）

为了让 Flatten 正常工作，我们需要给数据添加批次维度。可以通过以下方式实现：

batch_test_matrix = test_matrix[None, :]  # 添加额外维度表示批次
batch_test_matrix输出:
tensor([[[1, 2, 3],[4, 5, 6],[7, 8, 9]]])

• None：在第 0 维的位置插入一个新的维度
• :：保留原始张量的所有数据。

现在，数据已经包含批次维度，可以使用 Flatten 展开：

nn.Flatten()(batch_test_matrix)输出:
tensor([[1, 2, 3, 4, 5, 6, 7, 8, 9]])

将 Flatten 层加入 MLP

在 MLP 的构建中，Flatten 是第一步。我们将它加入模型的层列表中：

layers = [nn.Flatten()
]
layers输出:
[Flatten(start_dim=1, end_dim=-1)]

2.3.2 输入层

输入层连接展平后的图像到模型的其他部分。使用 nn.Linear 构建全连接层（densely connected），每个神经元及其权重会影响下一层的所有神经元。

layers = [nn.Flatten(),nn.Linear(input_size, 512),  # 输入层nn.ReLU()  # 输入层激活函数
]layers
输出:
[Flatten(start_dim=1, end_dim=-1),Linear(in_features=784, out_features=512, bias=True),ReLU()]

• 输入大小 input_size 为 1 x 28 x 28，即展平后为 784。
• 神经元数量设置为 512，可以通过调整值观察其对训练的影响。
• 使用 ReLU 作为激活函数以帮助网络捕获非线性特征。

2.3.3 隐藏层

增加一个隐藏层，进一步提取特征。隐藏层的神经元能够学习输入数据的特征表示，层数越多、神经元越多，模型就能提取更复杂、更抽象的特征。

• 第一隐藏层可能只学习简单的模式（如图像边缘或线条）。
• 第二隐藏层则可以在这些简单模式的基础上学习更高级的模式（如形状或局部结构）。

隐藏层是另一个全连接层，需要知道上一层的神经元数量作为输入大小。

layers = [nn.Flatten(),nn.Linear(input_size, 512),  # 输入层nn.ReLU(),  # 输入层激活函数nn.Linear(512, 512),  # 隐藏层nn.ReLU()  # 隐藏层激活函数
]

• 隐藏层神经元数量与输入层相同，均为 512。
• 使用 ReLU 激活函数。

2.3.4 输出层

输出层负责最终的分类预测。

n_classes = 10
layers = [nn.Flatten(),nn.Linear(input_size, 512),  # 输入层nn.ReLU(),  # 输入层激活函数nn.Linear(512, 512),  # 隐藏层nn.ReLU(),  # 隐藏层激活函数nn.Linear(512, n_classes)  # 输出层
]

• 不对输出层使用激活函数，而是通过损失函数处理模型的输出。
• 输出层的神经元数量等于分类的类别数（MNIST 数据集为 10）所以 n_classes=10，对应 10 个分类。

2.3.5 模型编译

model = nn.Sequential(*layers)
model输出:
Sequential((0): Flatten(start_dim=1, end_dim=-1)(1): Linear(in_features=784, out_features=512, bias=True)(2): ReLU()(3): Linear(in_features=512, out_features=512, bias=True)(4): ReLU()(5): Linear(in_features=512, out_features=10, bias=True)
)

• 使用 nn.Sequential 将所有层组合成一个顺序模型。
• 使用 *layers 解包列表，将层传递给 nn.Sequential。
  • 在 Python 中，* 是解包运算符，用于将一个可迭代对象（如列表、元组）中的元素依次取出，作为单独的参数传递给函数或构造器。

model.to(device)

• 将模型迁移到 GPU：默认情况下，模型在 CPU 上初始化。使用 .to(device) 方法将模型迁移到 GPU 上运行。

next(model.parameters()).device # 检查模型所在设备输出:
device(type='cuda', index=0)

• model.parameters() 返回的是模型所有参数（例如权重和偏置）的迭代器（generator 类型）。通过 next()，我们可以从迭代器中获取第一个参数（通常是第一个层的权重），然后通过 .device 属性查询它所在的设备

PyTorch 2.0 优化：
torch.compile 是 PyTorch 2.0 引入的新特性，用于动态编译和优化模型，旨在提升模型的执行效率。

model = torch.compile(model)

torch.compile 将模型包装为一个经过优化的模型对象，具体执行过程如下：

1. 捕获计算图：
   • 在模型的前向传播中，PyTorch 会捕获模型的计算图。
   • 计算图表示张量操作的顺序和依赖关系。
2. 编译优化：
   • PyTorch 使用后台优化工具（如 TorchDynamo 和 AOTAutograd）对计算图进行优化，包括：
     • 操作融合：将多个小的操作合并为一个大操作。
     • 内存优化：减少内存分配和回收的频率。
     • 内核优化：生成更高效的 GPU/CPU 内核代码。
3. 执行编译后的计算图：
   • 在模型训练或推理时，运行优化后的计算图，从而提升执行效率。

2.4 训练模型

现在我们可以使用训练数据训练模型，并用验证数据测试其性能。

2.4.1 损失函数与优化器

损失函数（Loss Function）：
模型需要通过一个“评分标准”来评估其预测的好坏。这里使用的是交叉熵损失函数（CrossEntropy），专门用于分类任务，评估模型对类别的预测是否准确。

loss_function = nn.CrossEntropyLoss()

优化器（Optimizer）：
优化器根据损失函数的评分（损失值）调整模型参数，从而逐渐提高模型的表现。这里使用 Adam 优化器：

optimizer = Adam(model.parameters())

2.4.2 计算准确率

虽然损失值能够反映模型的学习效果，但对人类而言很难直观理解，因此通常还会使用“准确率”来辅助评估模型性能。计算过程如下：

1. 比较模型预测值中每一批次的正确分类数与总样本数。
2. 使用如下函数计算每一批次的准确率：

def get_batch_accuracy(output, y, N):pred = output.argmax(dim=-1, keepdim=True)correct = pred.eq(y.view_as(pred)).sum().item()return correct / N

(1)pred = output.argmax(dim=-1, keepdim=True)

output是模型的前向传播结果，表示模型对每个样本的预测分数。它通常是一个二维张量，形状为 [batch_size, num_classes]。例子：

output = torch.tensor([[0.1, 0.5, 0.4],  # 第一个样本的分数[0.8, 0.1, 0.1]   # 第二个样本的分数
])

• 第一行 [0.1, 0.5, 0.4] 表示第一个样本的预测分数
  • 类别 0 的分数为 0.1，类别 1 的分数为 0.5，类别 2 的分数为 0.4 模型认为该样本最可能属于类别 1

再回来分析一下output.argmax的参数

• argmax(dim=-1)：沿着最后一个维度（num_classes）寻找分数最高的索引，即预测的类别。

• keepdim=True：保持输出张量的维度结构（即从 [batch_size, num_classes] 变为 [batch_size, 1]）

最终pred返回类别索引。

(2)correct = pred.eq(y.view_as(pred)).sum().item()

计算预测正确的样本数。

• y.view_as(pred)：将真实标签 y 的形状调整为与 pred 相同（从 [batch_size] 变为 [batch_size, 1]）。
• pred.eq(y.view_as(pred))：比较预测值 pred 和真实值 y，返回一个布尔张量，表示每个样本是否预测正确。
• sum()：对布尔张量求和，计算预测正确的样本数。
• .item()：将结果从张量转换为 Python 标量。

2.4.3 训练函数

定义一个train函数，用于对模型进行训练。其核心逻辑包括：

1. 初始化：将损失和准确率初始化为 0。
2. 训练循环：对于每个批次的数据：
   • 将数据加载到设备（如 GPU）；
   • 前向传播计算输出；
   • 使用损失函数计算损失；
   • 反向传播更新参数；
3. 记录结果：在每个批次中累计损失和准确率。

代码如下，具体见注释：

def train():# 初始化累计损失值为 0loss = 0# 初始化累计准确率为 0accuracy = 0# 将模型设置为训练模式，以启用 dropout 和 batch normalization 等训练特性model.train()# 遍历训练数据的每个批次for x, y in train_loader:# 将输入数据和标签移动到指定的设备（CPU 或 GPU）x, y = x.to(device), y.to(device)# 前向传播，使用模型对当前批次数据进行预测output = model(x)# 计算当前批次的损失值（如交叉熵损失）batch_loss = loss_function(output, y)# 清空优化器的梯度缓存，避免上次迭代的梯度影响当前计算optimizer.zero_grad()# 反向传播，计算损失函数对模型参数的梯度batch_loss.backward()# 使用优化器更新模型参数optimizer.step()# 累加当前批次的损失值，`item()` 将张量转换为标量loss += batch_loss.item()# 累加当前批次的准确率accuracy += get_batch_accuracy(output, y, train_N)# 打印训练的总损失值和准确率print(f"Train - Loss: {loss:.4f} Accuracy: {accuracy:.4f}")

这里详细地解释一下以下两个函数：

(1)model.train()

model.train() 是 PyTorch 中用于设置模型为训练模式的方法。这主要是为了让模型在训练时启用一些与训练相关的功能，例如：

1. 启用 Dropout：
   • 如果模型中包含 Dropout 层（用于随机丢弃神经元以防止过拟合），在训练模式下，Dropout 会随机丢弃一定比例的神经元。
   • 如果不调用 model.train()，Dropout 将默认禁用，这可能导致训练过程与实际推理不一致。
2. 启用 Batch Normalization 的动态更新：
   • 如果模型中包含 Batch Normalization 层，它会根据当前批次的数据统计均值和方差，并更新这些统计值。
   • 在训练模式下，Batch Normalization 层会动态计算和更新均值与方差。
   • 在验证或测试模式下（model.eval()），它会使用训练过程中计算的均值和方差。
3. 训练模式和评估模式的区别：
   • model.train()：用于训练阶段，启用 Dropout 和动态 Batch Normalization。
   • model.eval()：用于验证或测试阶段，禁用 Dropout 和动态 Batch Normalization。

(2)optimizer.zero_grad()

在 PyTorch 中，梯度是通过反向传播（backward()）计算的，每次调用 backward() 时，梯度会被累积到每个参数的 .grad 属性中。

• 累积是 PyTorch 的默认行为，梯度不会在每次反向传播后自动清除，而是将新计算的梯度累加到现有的梯度中。

在每次参数更新之前，调用 optimizer.zero_grad()，将所有参数的梯度清零，避免前一次计算的梯度影响当前的梯度更新。

为什么清零：

• 如果不清零，当前梯度会与上一轮的梯度累积，导致参数更新不准确。
• 通常，我们希望每次反向传播的梯度仅代表当前批次的贡献，而不是历史梯度的累积。

2.4.4 验证函数

validate() 函数用于在验证集上评估模型性能。

def validate():loss = 0accuracy = 0model.eval()with torch.no_grad():for x, y in valid_loader:x, y = x.to(device), y.to(device)output = model(x)loss += loss_function(output, y).item()accuracy += get_batch_accuracy(output, y, valid_N)print(f"Valid - Loss: {loss:.4f} Accuracy: {accuracy:.4f}")

与 train() 函数类似，但模型处于评估模式（model.eval()），且不需要更新参数（通过 torch.no_grad() 禁用梯度计算）。

with torch.no_grad()

torch.no_grad() 是一个临时的上下文管理器，表示在其作用范围内禁用梯度计算。这对于某些操作（如验证、推理等）非常重要，因为这些操作通常不需要反向传播或梯度更新。

2.4.5 训练循环

• 在训练和验证之间交替进行，观察模型的逐步改进。

• Epoch 的定义：完整遍历一次数据集称为一个 Epoch。

代码示例：训练 5 个 Epoch，并在每个 Epoch 后打印训练和验证的损失与准确率。

epochs = 5
for epoch in range(epochs):print(f"Epoch: {epoch + 1}")train()validate()输出:
Epoch: 0
Train - Loss: 56.6846 Accuracy: 0.9903
Valid - Loss: 25.8367 Accuracy: 0.9774
Epoch: 1
Train - Loss: 48.3223 Accuracy: 0.9917
Valid - Loss: 28.8761 Accuracy: 0.9776
Epoch: 2
Train - Loss: 37.2505 Accuracy: 0.9935
Valid - Loss: 32.4447 Accuracy: 0.9761
Epoch: 3
Train - Loss: 41.9876 Accuracy: 0.9931
Valid - Loss: 46.3217 Accuracy: 0.9727
Epoch: 4
Train - Loss: 36.6988 Accuracy: 0.9939
Valid - Loss: 30.7549 Accuracy: 0.9799

数据是在 DataLoader 中通过 shuffle=True 进行打乱的，当然还有可能因为Dropout等操作，导致了每次的结果都不太一样。通过多次 Epoch 和数据打乱，模型能逐步学习更稳定、更泛化的特征，有助于提升性能。

2.4.6 测试模型

可以将数据输入到训练好的模型中，得到输出预测值。

prediction = model(x_0_gpu)prediction
输出:
tensor([[-31.0694, -10.6213, -22.9587,   0.9323, -31.3773,  18.5830, -22.8076,-27.8728, -13.3324, -13.5257]], device='cuda:0',grad_fn=<AddmmBackward0>)

• 输出是 10 个数字（对应 10 个类别的预测分数）。
• 使用 argmax 找到分数最高的索引，即模型预测的类别。

prediction.argmax(dim=-1, keepdim=True)输出:
tensor([[5]], device='cuda:0')

再来看看实际的分类：

y_0输出:
5

说明最开始我们在2.2.1中显示的第一张图，手写数字5被正确识别。

3 总结

本文详细讲解了如何使用PyTorch构建多层感知器（MLP）模型，在经典的MNIST数据集上实现手写数字分类。文章从数据加载、预处理到模型搭建、训练和验证，逐步展示了完整的深度学习项目流程，同时结合代码深入解析关键技术点，如张量操作、激活函数、损失函数和优化器。