LeNet-5 在 MNIST 数据集上的 PyTorch 实现

LeNet-5 模型简介

LeNet-5 是由 Yann LeCun 等人在1998年提出的一种卷积神经网络结构,是深度学习领域的奠基性工作之一。尽管结构简单,但 LeNet-5 在手写数字识别等任务上取得了很好的效果,为后来更复杂的卷积神经网络奠定了基础。

LeNet-5 的网络结构如下:

LeNet-5 架构

主要包含:

2个卷积层
2个池化层
3个全连接层

这种结构能够有效地提取图像特征,并最终完成分类任务。

MNIST 数据集介绍

MNIST 是一个手写数字图像数据集,包含 60,000 张训练图像和 10,000 张测试图像。每张图像是 28x28 像素的灰度图,对应 0-9 的 10 个数字类别。

MNIST 数据集的一些样本如下:

MNIST 样本

由于数据简单且规模适中,MNIST 成为了机器学习和计算机视觉领域的经典入门数据集。

环境准备

首先我们需要安装必要的 Python 库:

pip install torch torchvision matplotlib

然后导入相关模块:

import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import datasets, transforms
import matplotlib.pyplot as plt

数据加载与预处理

使用 torchvision 可以方便地下载和加载 MNIST 数据集:

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

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)

这里我们对数据进行了归一化处理,使用的均值和标准差是 MNIST 数据集的统计值。

构建 LeNet-5 模型

接下来我们使用 PyTorch 定义 LeNet-5 网络结构:

class LeNet5(nn.Module):
    def __init__(self):
        super(LeNet5, self).__init__()
        self.conv1 = nn.Conv2d(1, 6, 5)
        self.conv2 = nn.Conv2d(6, 16, 5)
        self.fc1 = nn.Linear(16*4*4, 120)
        self.fc2 = nn.Linear(120, 84)
        self.fc3 = nn.Linear(84, 10)

    def forward(self, x):
        x = torch.relu(self.conv1(x))
        x = torch.max_pool2d(x, 2)
        x = torch.relu(self.conv2(x))
        x = torch.max_pool2d(x, 2)
        x = x.view(-1, 16*4*4)
        x = torch.relu(self.fc1(x))
        x = torch.relu(self.fc2(x))
        x = self.fc3(x)
        return x

model = LeNet5()

定义损失函数和优化器

我们使用交叉熵损失函数和 Adam 优化器:

criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)

模型训练

现在我们可以开始训练模型了:

num_epochs = 10
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model.to(device)

for epoch in range(num_epochs):
    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()
        
    print(f'Epoch {epoch+1}/{num_epochs}, Loss: {loss.item():.4f}')

模型评估

训练完成后,我们在测试集上评估模型性能:

model.eval()
correct = 0
total = 0
with torch.no_grad():
    for data, target in test_loader:
        data, target = data.to(device), target.to(device)
        outputs = model(data)
        _, predicted = torch.max(outputs.data, 1)
        total += target.size(0)
        correct += (predicted == target).sum().item()

print(f'Accuracy on test set: {100 * correct / total:.2f}%')

通过以上步骤,我们成功实现了 LeNet-5 模型,并在 MNIST 数据集上进行了训练和评估。一般情况下,该模型可以达到约 99% 的测试集准确率。

结果可视化

为了更直观地了解模型的表现,我们可以可视化一些预测结果:

def imshow(img):
    npimg = img.numpy()
    plt.imshow(np.transpose(npimg, (1, 2, 0)))
    plt.axis('off')

dataiter = iter(test_loader)
images, labels = next(dataiter)

outputs = model(images.to(device))
_, predicted = torch.max(outputs, 1)

fig = plt.figure(figsize=(25,4))
for idx in range(20):
    ax = fig.add_subplot(2, 10, idx+1, xticks=[], yticks=[])
    imshow(images[idx])
    ax.set_title(f"{predicted[idx]}", 
                 color=("green" if predicted[idx]==labels[idx] else "red"))
plt.show()

预测结果可视化