使用 PyTorch 构建 MNIST 手写数字识别模型

引言

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 张量。
  1. 展示部分手写数字图像

    python 复制代码
    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()

    解释:

    • 使用 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

解释:

  • 训练模型并计算损失。
  • 使用梯度下降更新权重。
  1. 定义测试函数

    python 复制代码
    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):>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 加速训练过程。希望这篇博客对你有所帮助!

相关推荐
java1234_小锋36 分钟前
TensorFlow2 Python深度学习 - TensorFlow2框架入门 - 计算图和 tf.function 简介
python·深度学习·tensorflow·tensorflow2
程序员晚枫39 分钟前
Python 3.14新特性:Zstandard压缩库正式加入标准库,性能提升30%
python
szxinmai主板定制专家43 分钟前
基于 ZYNQ ARM+FPGA+AI YOLOV4 的电网悬垂绝缘子缺陷检测系统的研究
arm开发·人工智能·嵌入式硬件·yolo·fpga开发
逆境清醒43 分钟前
VS Code配置Python开发环境系列(1)___VScode的安装 ,VScode常用快捷键
vscode·python·visual studio code
万粉变现经纪人1 小时前
如何解决 pip install -r requirements.txt 无效可编辑项 ‘e .‘(-e 拼写错误)问题
开发语言·python·r语言·beautifulsoup·pandas·pip·scipy
聚客AI1 小时前
🌈提示工程已过时?上下文工程从理论到实践的完整路线图
人工智能·llm·agent
C嘎嘎嵌入式开发1 小时前
(二) 机器学习之卷积神经网络
人工智能·机器学习·cnn
潇凝子潇1 小时前
在使用Nacos作为注册中心和配置中心时,如何解决服务发现延迟或配置更新不及时的问题
开发语言·python·服务发现
烛阴1 小时前
Python 列表推导式:让你的代码更优雅、更高效
前端·python
文心快码BaiduComate1 小时前
开工不累,双强护航:文心快码接入 DeepSeek-V3.2-Exp和 GLM-4.6,助你节后高效Coding
前端·人工智能·后端