一、前言

MNIST 手写数字数据集是深度学习领域的 “Hello World”，非常适合入门者学习深度学习模型的构建与训练。本文将结合 PyTorch 框架，带大家完整走一遍从数据准备到模型训练、评估的流程。

二、数据准备

（一）导入工具库

首先，我们需要导入一系列必要的库，包括用于数值计算的numpy，PyTorch 相关的模块（用于构建模型、数据处理、优化等），以及用于数据可视化的matplotlib。

python

运行

import numpy as np
import torch
from torchvision.datasets import MNIST
import torchvision.transforms as transforms
from torch.utils.data import DataLoader
import torch.nn.functional as F
import torch.optim as optim
from torch import nn
from torch.utils.tensorboard import SummaryWriter
import matplotlib.pyplot as plt
%matplotlib inline

（二）定义超参数

超参数会影响模型的训练过程和性能，这里定义了批次大小、学习率和训练轮数。

python

运行

train_batch_size = 64
test_batch_size = 128
learning_rate = 0.01
num_epochs = 20

（三）数据预处理与加载

使用transforms对数据进行预处理，将图像转换为张量并进行归一化。然后利用MNIST数据集类下载并加载数据，通过DataLoader创建数据迭代器，方便批量处理数据。

python

运行

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

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

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

（四）数据可视化

为了直观了解数据，我们从测试集中取出一批数据并可视化其中的部分图像及其真实标签。

python

运行

examples = enumerate(test_loader)
batch_idx, (example_data, example_targets) = next(examples)

fig = plt.figure()
for i in range(6):
    plt.subplot(2, 3, i + 1)
    plt.tight_layout()
    plt.imshow(example_data[i][0], cmap='gray', interpolation='none')
    plt.title("Ground Truth: {}".format(example_targets[i]))
    plt.xticks([])
    plt.yticks([])

运行这段代码，会看到 6 张手写数字图像，每张图像上方标注了其真实的数字标签，能让我们对 MNIST 数据有更直观的认识。

三、构建模型

我们定义一个包含两个隐藏层的神经网络模型，使用ReLU作为激活函数，输出层使用Softmax函数将输出转换为类别概率分布。

python

运行

class Net(nn.Module):
    def __init__(self, in_dim, n_hidden_1, n_hidden_2, out_dim):
        super(Net, self).__init__()
        self.flatten = nn.Flatten()
        self.layer1 = nn.Sequential(nn.Linear(in_dim, n_hidden_1), nn.BatchNorm1d(n_hidden_1))
        self.layer2 = nn.Sequential(nn.Linear(n_hidden_1, n_hidden_2), nn.BatchNorm1d(n_hidden_2))
        self.out = nn.Sequential(nn.Linear(n_hidden_2, out_dim))

    def forward(self, x):
        x = self.flatten(x)
        x = F.relu(self.layer1(x))
        x = F.relu(self.layer2(x))
        x = F.softmax(self.out(x), dim=1)
        return x

这里，nn.Flatten用于将输入的二维图像（28×28）展平为一维向量，方便全连接层处理；BatchNorm1d用于对隐藏层的输出进行批量归一化，有助于加速模型训练和提升性能；ReLU激活函数引入非线性，增强模型的表达能力；Softmax函数则将最后一层的输出转换为各个类别的概率。

四、模型训练与评估

（一）实例化模型、定义损失函数和优化器

根据硬件情况选择使用 GPU 或 CPU，实例化模型并将其移动到相应设备上。损失函数选择交叉熵损失，优化器使用带动量的随机梯度下降（SGD）。

python

运行

lr = 0.01
momentum = 0.9

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
model = Net(28 * 28, 300, 100, 10)
model.to(device)

criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(), lr=lr, momentum=momentum)

（二）训练模型

训练过程包括多轮迭代（epoch），每轮迭代中遍历训练数据，进行前向传播、计算损失、反向传播和参数更新，同时记录训练损失和准确率。每轮结束后，在测试集上评估模型性能。还采用了动态调整学习率的策略，在训练到一定轮数后减小学习率，以提高模型收敛精度。

python

运行

losses = []
acces = []
eval_losses = []
eval_acces = []
writer = SummaryWriter(log_dir='logs', comment='train-loss')

for epoch in range(num_epochs):
    train_loss = 0
    train_acc = 0
    model.train()
    if epoch % 5 == 0:
        optimizer.param_groups[0]['lr'] *= 0.9
        print("学习率:{:.6f}".format(optimizer.param_groups[0]['lr']))
    for img, label in train_loader:
        img = img.to(device)
        label = label.to(device)
        out = model(img)
        loss = criterion(out, label)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        train_loss += loss.item()
        writer.add_scalar('Train', train_loss / len(train_loader), epoch)
        _, pred = out.max(1)
        num_correct = (pred == label).sum().item()
        acc = num_correct / img.shape[0]
        train_acc += acc

    losses.append(train_loss / len(train_loader))
    acces.append(train_acc / len(train_loader))
    eval_loss = 0
    eval_acc = 0
    model.eval()
    for img, label in test_loader:
        img = img.to(device)
        label = label.to(device)
        img = img.view(img.size(0), -1)
        out = model(img)
        loss = criterion(out, label)
        eval_loss += loss.item()
        _, pred = out.max(1)
        num_correct = (pred == label).sum().item()
        acc = num_correct / img.shape[0]
        eval_acc += acc

    eval_losses.append(eval_loss / len(test_loader))
    eval_acces.append(eval_acc / len(test_loader))
    print('epoch: {}, Train Loss: {:.4f}, Train Acc: {:.4f}, Test Loss: {:.4f}, Test Acc: {:.4f}'
          .format(epoch, train_loss / len(train_loader), train_acc / len(train_loader),
                  eval_loss / len(test_loader), eval_acc / len(test_loader)))

在训练过程中，model.train()将模型设置为训练模式，启用 dropout 和 batch normalization 的训练行为；model.eval()则将模型设置为评估模式，关闭 dropout 等，使用 batch normalization 的已学参数。动态调整学习率是因为在训练初期，较大的学习率可以加快模型收敛速度，而到了训练后期，减小学习率有助于模型更精确地收敛到最优解。

（三）可视化训练损失

训练完成后，我们可以可视化训练损失的变化趋势，直观地观察模型的训练过程。

python

运行

plt.title('train loss')
plt.plot(np.arange(len(losses)), losses)
plt.legend(['Train Loss'], loc='upper right')

运行这段代码，会得到一张训练损失随训练轮数变化的折线图，能看到损失逐渐下降，说明模型在不断学习和优化。

五、总结

通过本次基于 PyTorch 的 MNIST 手写数字识别实战，我们完整地经历了深度学习项目的主要流程：数据准备与预处理、模型构建、模型训练与评估。在这个过程中，我们还应用了一些技巧，比如动态调整学习率来优化模型训练，使用批量归一化提升模型性能等。MNIST 作为入门数据集，让我们很好地理解了深度学习的基本概念和方法，为后续处理更复杂的深度学习任务打下了基础。