Pytorch 实战系列 9——MNIST 实战

2023-11-18 10:51:33

前言

在本文中，我们将使用 Pytorch 实现一个 MNIST 图像分类器。MNIST 数据集是一个由手写数字图像组成的大型数据集，广泛用于机器学习和深度学习的图像分类任务。通过本教程，您将学习如何使用 Pytorch 构建和训练神经网络模型，并使用 MNIST 数据集对其进行评估。

MNIST 数据集简介

MNIST 数据集包含 70,000 张手写数字图像，其中 60,000 张用于训练，10,000 张用于测试。每幅图像都是 28x28 像素的灰度图像，代表一个数字（从 0 到 9）。

MNIST 数据集是图像分类任务的基准，广泛用于评估机器学习和深度学习模型的性能。

Pytorch简介

Pytorch 是一个流行的 Python 深度学习库，它提供了构建和训练神经网络模型所需的工具和功能。Pytorch 以其灵活性、易用性和对大规模并行训练的支持而闻名。

在本文中，我们将使用 Pytorch 的以下组件：

torch.nn：神经网络模型构建模块
torch.optim：优化算法
torch.utils.data：数据加载器

构建神经网络模型

我们将使用卷积神经网络（CNN）模型来对 MNIST 图像进行分类。CNN 是专门为处理图像数据而设计的深度学习模型。

我们的 CNN 模型将包含以下层：

卷积层：使用过滤器提取图像中的特征。
池化层：对特征图进行下采样以减少维度。
全连接层：将提取的特征映射到数字预测。

训练神经网络模型

一旦我们构建了神经网络模型，我们就需要使用 MNIST 数据集对其进行训练。训练过程涉及以下步骤：

数据加载： 使用 Pytorch 数据加载器将 MNIST 数据集加载到内存中。
模型初始化： 初始化神经网络模型的权重和偏差。
前向传播： 将图像输入到模型并计算输出预测。
损失计算： 计算模型预测与真实标签之间的损失。
反向传播： 使用反向传播算法计算损失函数的梯度。
权重更新： 使用优化算法更新模型的权重和偏差。

评估神经网络模型

在训练神经网络模型后，我们需要使用测试集对其进行评估。评估过程涉及以下步骤：

加载测试集： 使用 Pytorch 数据加载器将 MNIST 测试集加载到内存中。
模型预测： 将测试图像输入到模型并获得预测。
准确率计算： 将模型预测与真实标签进行比较以计算准确率。

完整代码

本节提供了 Pytorch MNIST 实战的完整代码。该代码涵盖了数据加载、模型构建、训练和评估的所有步骤。

# 导入必要的库
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader

# 定义神经网络模型
class CNN(nn.Module):
    def __init__(self):
        super(CNN, self).__init__()
        self.conv1 = nn.Conv2d(1, 32, 3, 1)
        self.pool1 = nn.MaxPool2d(2, 2)
        self.conv2 = nn.Conv2d(32, 64, 3, 1)
        self.pool2 = nn.MaxPool2d(2, 2)
        self.fc1 = nn.Linear(64 * 4 * 4, 128)
        self.fc2 = nn.Linear(128, 10)

    def forward(self, x):
        x = self.pool1(F.relu(self.conv1(x)))
        x = self.pool2(F.relu(self.conv2(x)))
        x = x.view(x.size(0), -1)
        x = F.relu(self.fc1(x))
        x = self.fc2(x)
        return x

# 加载 MNIST 数据集
train_dataset = torchvision.datasets.MNIST(root='./data', train=True, download=True, transform=transforms.ToTensor())
train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True)

test_dataset = torchvision.datasets.MNIST(root='./data', train=False, download=True, transform=transforms.ToTensor())
test_loader = DataLoader(test_dataset, batch_size=64, shuffle=False)

# 构建神经网络模型
model = CNN()

# 定义损失函数和优化器
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(), lr=0.01)

# 训练神经网络模型
for epoch in range(10):
    running_loss = 0.0
    for i, data in enumerate(train_loader, 0):
        # 获取输入和标签
        inputs, labels = data

        # 前向传播
        outputs = model(inputs)

        # 损失计算
        loss = criterion(outputs, labels)

        # 反向传播和权重更新
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        # 打印训练信息
        running_loss += loss.item()
        if i % 2000 == 1999:    # 每 2000 个小批量打印一次
            print('[%d, %5d] loss: %.3f' %
                  (epoch + 1, i + 1, running_loss / 2000))
            running_loss = 0.0

# 评估神经网络模型
correct = 0
total = 0
with torch.no_grad():
    for data in test_loader:
        images, labels = data
        outputs = model(images)
        _, predicted = torch.max(outputs.data, 1)
        total += labels.size(0)
        correct += (predicted == labels).sum().item()

print('准确率: %d %%' % (100 * correct / total))