手写数字识别是计算机视觉领域的经典任务，而 MNIST 数据集提供了一个标准化的基准。本文将使用 PyTorch 框架构建一个标准的卷积神经网络（CNN），对 MNIST 数据集进行分类，并展示完整的训练和测试流程。
1. 环境准备与数据加载
在深度学习中，数据预处理是至关重要的步骤。MNIST 数据集包含 28×28 灰度图像，训练集 60000 张，测试集 10000 张。我们使用 torchvision 提供的工具来加载数据，并进行标准化处理：
import torch
from torchvision import datasets, transforms
# 数据预处理：转换为 Tensor 并标准化到 [-1, 1]
transform = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.5,), (0.5,))
])
# 加载训练集和测试集
train_dataset = datasets.MNIST(root='./data', train=True, download=True, transform=transform)
test_dataset = datasets.MNIST(root='./data', train=False, download=True, transform=transform)
# 数据加载器
train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=64, shuffle=True)
test_loader = torch.utils.data.DataLoader(test_dataset, batch_size=1000, shuffle=False)说明：
ToTensor()会将图像从[0, 255]转换为[0, 1]的浮点数张量。
Normalize(mean, std)将张量标准化为(x - mean) / std)，有助于加快训练收敛。
2. 构建卷积神经网络
我们定义一个两层卷积的 CNN 模型，结构如下：
-
卷积层1：输入 1 通道 → 输出 32 通道，卷积核 3×3
-
卷积层2：输入 32 通道 → 输出 64 通道，卷积核 3×3
-
池化层：使用 2×2 最大池化
-
全连接层1：将特征映射展平后映射到 128 个神经元
-
输出层：映射到 10 类数字
import torch.nn as nn
import torch.nn.functional as Fclass CNN(nn.Module):
def init(self):
super(CNN, self).init()
self.conv1 = nn.Conv2d(1, 32, 3, padding=1)
self.conv2 = nn.Conv2d(32, 64, 3, padding=1)
self.pool = nn.MaxPool2d(2, 2)
self.fc1 = nn.Linear(64 * 7 * 7, 128)
self.fc2 = nn.Linear(128, 10)def forward(self, x): x = F.relu(self.conv1(x)) x = self.pool(x) x = F.relu(self.conv2(x)) x = self.pool(x) x = x.view(x.size(0), -1) x = F.relu(self.fc1(x)) x = self.fc2(x) return x
说明：
view(x.size(0), -1)将二维特征图展平为向量，便于全连接层处理。
ReLU激活函数增加非线性能力。
3. 定义训练环境
在训练前，需要将模型移动到 GPU（若可用），并定义损失函数与优化器：
import torch.optim as optim
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = CNN().to(device)
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)说明：
CrossEntropyLoss适用于多分类任务。
Adam优化器在大多数场景下表现稳定且收敛速度快。
4. 训练函数
训练函数包括以下步骤：
-
将模型设为训练模式
model.train() -
遍历训练数据
-
前向传播 → 计算损失 → 反向传播 → 更新参数
-
每 100 个 batch 输出一次训练信息
def train(model, device, train_loader, optimizer, criterion, epoch):
model.train()
for batch_idx, (data, target) in enumerate(train_loader):
data, target = data.to(device), target.to(device)
optimizer.zero_grad()
output = model(data)
loss = criterion(output, target)
loss.backward()
optimizer.step()if batch_idx % 100 == 0: print(f'Train Epoch: {epoch} [{batch_idx * len(data)}/{len(train_loader.dataset)} ' f'({100. * batch_idx / len(train_loader):.0f}%)]\tLoss: {loss.item():.6f}')
5. 测试函数
测试阶段不更新模型参数，只计算平均损失和准确率：
def test(model, device, test_loader, criterion):
model.eval()
test_loss = 0
correct = 0
with torch.no_grad():
for data, target in test_loader:
data, target = data.to(device), target.to(device)
output = model(data)
test_loss += criterion(output, target).item()
pred = output.argmax(dim=1, keepdim=True)
correct += pred.eq(target.view_as(pred)).sum().item()
test_loss /= len(test_loader)
accuracy = 100. * correct / len(test_loader.dataset)
print(f'\nTest set: Average loss: {test_loss:.4f}, Accuracy: {correct}/{len(test_loader.dataset)} '
f'({accuracy:.2f}%)\n')说明：
使用
torch.no_grad()禁止梯度计算，提高测试效率并减少显存占用。
argmax获取预测类别。
6. 执行训练与测试
设定训练轮数 num_epochs，依次调用训练和测试函数：
num_epochs = 5
for epoch in range(1, num_epochs + 1):
train(model, device, train_loader, optimizer, criterion, epoch)
test(model, device, test_loader, criterion)训练与测试结果示例：
Train Epoch: 1 [0/60000 (0%)] Loss: 0.384921
...
Test set: Average loss: 0.0543, Accuracy: 9812/10000 (98.12%)7. 总结
本文使用 PyTorch 实现了一个标准 CNN 模型完成 MNIST 手写数字分类任务，模型结构简单、易于理解，适合作为入门和实验。经过 5 个训练周期，模型可以达到 约 98% 的测试准确率，显示出卷积神经网络在图像分类中的有效性。
后续可以进一步优化：
-
增加卷积层或残差结构
-
使用数据增强提升泛化能力
-
尝试学习率调度器（LR scheduler）或其他优化器