用卷积神经网络 (CNN) 实现 MNIST 手写数字识别-CSDN博客

本文链接：https://blog.csdn.net/ycy200377/article/details/147701218

在深度学习领域，MNIST 手写数字识别是经典的入门级项目，就像编程世界里的 “Hello, World”。卷积神经网络（Convolutional Neural Network，CNN）作为处理图像数据的强大工具，在该任务中展现出卓越的性能。本文将结合具体的 PyTorch 代码，详细解析如何利用 CNN 实现 MNIST 手写数字识别，带大家从代码实践深入理解背后的技术原理。

一、数据准备：加载与预处理 MNIST 数据集

MNIST 数据集包含 6 万张训练图像和 1 万张测试图像，涵盖 0 - 9 这十个数字的手写体。我们借助torchvision库中的datasets.MNIST函数来加载数据，具体代码如下：

import torch
from torch import nn
from torch.utils.data import DataLoader
from torchvision import datasets
from torchvision.transforms import ToTensor

training_data = datasets.MNIST(
    root="data",
    train=True,
    download=True,
    transform=ToTensor(),
)
test_data = datasets.MNIST(
    root="data",
    train=False,
    download=True,
    transform=ToTensor(),
)

上述代码中，root="data"指定数据集的存储路径；train=True表示加载训练集，train=False用于加载测试集；download=True确保本地无数据集时自动下载；transform=ToTensor()将图像数据转换为 PyTorch 张量格式，并把像素值从 0 - 255 归一化到 0 - 1 区间，便于后续处理。

为直观感受数据，我们用matplotlib库绘制 9 张训练图像及其标签：

from matplotlib import pyplot as plt
figure = plt.figure()
for i in range(9):
    img, label = training_data[i + 59000]
    figure.add_subplot(3, 3, i + 1)
    plt.title(label)
    plt.axis("off")
    plt.imshow(img.squeeze(), cmap="gray")
    a = img.squeeze()
plt.show()

完成数据加载后，使用DataLoader将数据封装成批次，方便模型训练和测试：

train_dataloader = DataLoader(training_data, batch_size=64)
test_dataloader = DataLoader(test_data, batch_size=64)

batch_size=64意味着每次训练或测试，模型会同时处理 64 个样本，能提高计算效率和训练稳定性。

二、模型构建：搭建卷积神经网络架构

我们定义一个名为CNN的类，继承自nn.Module，用于构建卷积神经网络：

class CNN(nn.Module):
    def __init__(self):
        super(CNN, self).__init__()
        self.conv1 = nn.Sequential(
            nn.Conv2d(
                in_channels=1,
                out_channels=16,
                kernel_size=3,
                stride=1,
                padding=1,
            ),
            nn.ReLU(),
            nn.MaxPool2d(2)
        )
        self.conv2 = nn.Sequential(
            nn.Conv2d(16, 32, 3, 1, 1),
            nn.ReLU(),
            nn.MaxPool2d(2),
        )
        self.conv3 = nn.Sequential(
            nn.Conv2d(32, 64, 3, 1, 1),
            nn.ReLU(),
        )
        self.out = nn.Linear(64 * 7 * 7, 10)

    def forward(self, x):
        x = self.conv1(x)
        x = self.conv2(x)
        x = self.conv3(x)
        x = x.view(x.size(0), -1)
        output = self.out(x)
        return output

卷积层（nn.Conv2d）：在conv1、conv2和conv3中，通过卷积层提取图像特征。例如conv1中的nn.Conv2d(in_channels=1, out_channels=16, kernel_size=3, stride=1, padding=1)，in_channels=1表示输入图像为单通道灰度图，out_channels=16表示输出 16 个特征图，kernel_size=3指定 3×3 的卷积核，stride=1是步长，padding=1用于保持图像尺寸不变。
激活函数（nn.ReLU）：紧跟在卷积层之后，为模型引入非线性，帮助模型学习复杂的模式。
池化层（nn.MaxPool2d）：通过下采样操作，如nn.MaxPool2d(2)将图像尺寸减半，减少数据量和模型参数，同时保留重要特征，防止过拟合。
全连接层（nn.Linear）：self.out = nn.Linear(64 * 7 * 7, 10)将卷积层输出的特征图展平后连接到全连接层，输出 10 个神经元对应 0 - 9 十个数字类别，完成最终分类。

最后，将模型移动到合适的计算设备（GPU、MPS 或 CPU）上：

device = "cuda" if torch.cuda.is_available() else "mps" if torch.backends.mps.is_available() else "cpu"
print(f"Using {device} device")
model = CNN().to(device)
print(model)

三、模型训练与测试：优化与评估

3.1 训练函数

def train(dataloader, model, loss_fn, optimizer):
    model.train()
    batch_size_num = 1
    for X, y in dataloader:
        X, y = X.to(device), y.to(device)
        pred = model.forward(X)
        loss = loss_fn(pred, y)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        loss_value = loss.item()
        if batch_size_num % 100 == 0:
            print(f"loss: {loss_value:>7f} [number:{batch_size_num}]")
        batch_size_num += 1

在训练函数中，model.train()将模型设为训练模式。遍历数据加载器，将每一批数据和标签移至指定设备，前向传播计算预测值，通过交叉熵损失函数nn.CrossEntropyLoss()计算损失，optimizer.zero_grad()清空梯度，loss.backward()反向传播计算梯度，optimizer.step()更新模型参数，每 100 个批次打印一次损失值。

3.2 测试函数

def test(dataloader, model, loss_fn):
    size = len(dataloader.dataset)
    num_batches = len(dataloader)
    model.eval()
    test_loss, correct = 0, 0
    with torch.no_grad():
        for X, y in dataloader:
            pred = model(X)
            test_loss += loss_fn(pred, y).item()
            correct += (pred.argmax(1) == y).type(torch.float).sum().item()
    test_loss /= num_batches
    correct /= size
    print(f"Test Error: \n Accuracy: {(100 * correct):>0.1f}%, Avg loss: {test_loss:>8f} \n")
    return test_loss, correct

测试函数中，model.eval()将模型设为评估模式，关闭如 Dropout 等训练时的操作。在with torch.no_grad()下遍历测试数据，计算测试损失和正确预测的样本数，最后计算平均损失和准确率并输出。

3.3 执行训练与测试

loss_fn = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=0.01)
epochs = 10
for t in range(epochs):
    print(f"Epoch {t + 1}\n--------------------")
    train(train_dataloader, model, loss_fn, optimizer)
print("Done!")
test(test_dataloader, model, loss_fn)

我们选用交叉熵损失函数和 Adam 优化器，学习率设为 0.01，通过 10 个训练周期不断优化模型，训练完成后在测试集上评估模型性能，得到最终的准确率和平均损失。