PyTorch探索之MNIST数据集解析与识别

人工智能

2023-11-02 15:22:30

在人工智能领域，PyTorch是一个新星，它是一种基于Python的深度学习框架，因其简单、灵活和强大的功能，受到了越来越多研究人员和工程师的青睐。

为了帮助大家更好地理解和使用PyTorch，我们以一个具体的例子——MNIST数据集的手写数字识别，来演示PyTorch的强大功能。

MNIST数据集

MNIST数据集是一个非常经典的手写数字数据集，它包含70,000张手写数字图片，其中60,000张用于训练，10,000张用于测试。每张图片都是一个28×28的灰度图像，代表一个手写数字。

PyTorch构建卷积神经网络

为了识别MNIST数据集中的手写数字，我们将使用PyTorch构建一个卷积神经网络（CNN）。CNN是一种非常适合图像识别的深度学习模型，它能够自动学习图像中的特征，并将其提取出来，以便进行分类。

PyTorch的CNN模型结构如下：

输入层：该层接收MNIST数据集中的手写数字图片。
卷积层：该层使用一组可训练的滤波器来提取图像中的特征。
池化层：该层将卷积层提取的特征进行降维，以减少模型的参数数量。
全连接层：该层将池化层提取的特征转换为最终的分类结果。

训练和测试模型

我们将使用PyTorch的训练和测试函数来训练和测试CNN模型。训练函数将使用MNIST数据集中的训练数据来训练模型，测试函数将使用MNIST数据集中的测试数据来测试模型的性能。

训练和测试模型的代码如下：

import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import datasets, transforms

# 定义CNN模型
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.conv1(x)
        x = self.pool1(x)
        x = self.conv2(x)
        x = self.pool2(x)
        x = x.view(-1, 64 * 4 * 4)
        x = self.fc1(x)
        x = self.fc2(x)
        return x

# 实例化CNN模型
model = CNN()

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

# 训练模型
for epoch in range(10):
    # 遍历训练集
    for i, (inputs, labels) in enumerate(train_loader):
        # 将输入数据和标签转换为PyTorch张量
        inputs = inputs.to(device)
        labels = labels.to(device)

        # 前向传播
        outputs = model(inputs)

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

        # 反向传播
        loss.backward()

        # 更新模型参数
        optimizer.step()

        # 打印训练信息
        if i % 100 == 0:
            print(f'Epoch: {epoch + 1}, Iteration: {i + 1}, Loss: {loss.item()}')

# 测试模型
model.eval()
with torch.no_grad():
    # 遍历测试集
    for i, (inputs, labels) in enumerate(test_loader):
        # 将输入数据和标签转换为PyTorch张量
        inputs = inputs.to(device)
        labels = labels.to(device)

        # 前向传播
        outputs = model(inputs)

        # 计算准确率
        _, predicted = torch.max(outputs.data, 1)
        accuracy = (predicted == labels).sum().item() / len(labels)

        # 打印测试信息
        print(f'Test Accuracy: {accuracy * 100}%')