引言
MNIST 数据集是一个经典的机器学习数据集,包含了 70,000 张手写数字图像,其中 60,000 张用于训练,10,000 张用于测试。每张图像都是 28x28 像素的灰度图像,并且已经被居中处理以减少预处理步骤。本文将介绍如何使用 PyTorch 构建一个简单的神经网络来识别 MNIST 数据集中的手写数字。
环境准备
首先确保你已经安装了 PyTorch 库。可以通过以下命令安装:
            
            
              python
              
              
            
          
          pip install torch torchvision数据加载与预处理
导入必要的库
            
            
              python
              
              
            
          
          import torch
from torch import nn
from torch.utils.data import DataLoader
from torchvision import datasets
from torchvision.transforms import ToTensor下载训练数据集
            
            
              python
              
              
            
          
          training_data = datasets.MNIST(
    root="data",
    train=True,
    download=True,
    transform=ToTensor()
)下载测试数据集
            
            
              python
              
              
            
          
          test_data = datasets.MNIST(
    root="data",
    train=False,
    download=True,
    transform=ToTensor()
)解释:
- 使用 datasets.MNIST下载 MNIST 数据集,并将其保存在"data"目录下。
- train=True表示下载训练集,- train=False表示下载测试集。
- download=True表示如果数据集不存在,则下载数据集。
- transform=ToTensor()将图像数据转换为 PyTorch 张量。
- 
展示部分手写数字图像 pythonfrom 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()解释: - 使用 matplotlib库展示一些手写数字图像。
 
- 使用 
数据加载器
创建数据加载器
            
            
              python
              
              
            
          
          train_dataloader = DataLoader(training_data, batch_size=64)
test_dataloader = DataLoader(test_data, batch_size=64)解释:
- 使用 DataLoader将数据打包成批次,batch_size=64表示每次迭代返回 64 个样本。
检查数据形状
            
            
              python
              
              
            
          
          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解释:
- 查看数据的形状,确保数据已经被正确打包。
模型定义
定义神经网络模型
            
            
              python
              
              
            
          
          class NeuralNetwork(nn.Module):
    def __init__(self):
        super().__init__()
        self.flatten = nn.Flatten()
        self.hidden1 = nn.Linear(28 * 28, 128)
        self.hidden2 = nn.Linear(128, 256)
        self.hidden3 = nn.Linear(256, 128)
        self.hidden4 = nn.Linear(128, 256)
        self.hidden5 = nn.Linear(256, 128)
        self.out = nn.Linear(128, 10)
    def forward(self, x):
        x = self.flatten(x)
        x = self.hidden1(x)
        x = torch.sigmoid(x)
        x = self.hidden2(x)
        x = torch.sigmoid(x)
        x = self.hidden3(x)
        x = torch.sigmoid(x)
        x = self.hidden4(x)
        x = torch.sigmoid(x)
        x = self.hidden5(x)
        x = torch.sigmoid(x)
        x = self.out(x)
        return x解释:
- 定义一个简单的多层感知器模型。
- 使用 nn.Flatten展平输入图像。
- 多个全连接层 (nn.Linear) 和激活函数 (torch.sigmoid)。
实例化模型并移动到 GPU
            
            
              python
              
              
            
          
          device = "cuda" if torch.cuda.is_available() else "mps" if torch.backends.mps.is_available() else 'cpu'
print(f"Using {device} device")
model = NeuralNetwork().to(device)
print(model)解释:
- 检查是否有可用的 GPU 并将模型移动到 GPU 上。
训练与测试
定义损失函数和优化器
            
            
              python
              
              
            
          
          loss_fn = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=0.01)定义训练函数
            
            
              python
              
              
            
          
          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解释:
- 训练模型并计算损失。
- 使用梯度下降更新权重。
- 
定义测试函数 pythondef 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):>0.1f}%, Avg loss: {test_loss:>8f}")解释: - 测试模型并计算准确率和平均损失。
 
训练和测试模型
            
            
              python
              
              
            
          
          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)解释:
- 进行多次训练周期(epoch)。
- 最后测试模型。
结果













总结
通过上述步骤,我们构建了一个简单的神经网络模型来识别 MNIST 数据集中的手写数字。使用 PyTorch 可以轻松地处理数据、定义模型、训练和测试模型。这个项目不仅展示了如何使用 PyTorch 进行图像识别,还介绍了如何利用 GPU 加速训练过程。希望这篇博客对你有所帮助!