卷积神经网络实战：手写数字识别

1. 引言

在机器学习和深度学习领域，手写数字识别是一个经典的应用案例。通过训练模型来识别手写数字，不仅可以帮助我们理解神经网络的基本原理，还能让我们接触到实际的图像处理任务。本篇博客将介绍如何使用卷积神经网络（Convolutional Neural Networks, CNN）来识别MNIST数据集中的手写数字。

2. 数据准备

MNIST数据集包含了70,000张手写数字图像，其中60,000张用于训练，10,000张用于测试。这些图像是灰度的，大小为28x28像素，并且已经被居中处理，以简化预处理步骤并加快运行速度。

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(),
)

print(len(training_data))

3. 可视化数据

为了更好地理解数据集的内容，我们可以可视化部分图像。

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")

plt.show()

4. 创建数据加载器

数据加载器用于打包数据，并允许我们批量处理数据。

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

for X, y in test_dataloader:
    print(f"Shape of X [N,C,H,W]: {X.shape}")
    print(f"Shape of y: {y.shape} {y.dtype}")
    break

5. 设备选择

检查当前设备是否支持GPU加速，这对于加快训练过程非常重要。

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

6. 构建卷积神经网络

接下来定义一个简单的CNN模型，用于识别手写数字。

class CNN(nn.Module):
    def __init__(self):
        super(CNN, self).__init__()
        self.conv1 = nn.Sequential(
            nn.Conv2d(1, 16, 5, 1, 2),
            nn.ReLU(),
            nn.MaxPool2d(2),
        )
        self.conv2 = nn.Sequential(
            nn.Conv2d(16, 32, 5, 1, 2),
            nn.ReLU(),
            nn.MaxPool2d(2),
        )
        self.out = nn.Linear(32 * 7 * 7, 10)

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

model = CNN().to(device)
print(model)

7. 训练与测试

定义训练和测试函数，使用Adam优化器和交叉熵损失函数进行训练。

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(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

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:
            X, y = X.to(device), y.to(device)
            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 result:\n Accuracy: {(100 * correct)}%, Avg loss: {test_loss}")

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)

输出结果：

8. 结语

通过以上步骤，我们成功构建并训练了一个卷积神经网络来识别MNIST数据集中的手写数字。这个过程不仅加深了我们对CNN的理解，还为我们提供了处理图像数据的实践经验。未来，我们可以尝试不同的网络结构或者调整超参数来进一步优化模型的表现。