手把手教你实现PyTorch的MNIST数据集

1、什么是MNIST数据集

MNIST数据集为手写数字图像数据集，包含60,000张训练图像和10,000张测试图像。每个图像大小为28x28像素，标注为0到9的数字。MNIST数据集可以被广泛地应用于图像识别、机器学习等相关领域。

2、下载MNIST数据集

2.1、使用PyTorch自带的函数下载

在PyTorch中，可以使用以下代码下载MNIST数据集。

import torch

from torchvision import datasets
from torchvision.transforms import ToTensor
# 创建MNIST数据集
train_data = datasets.MNIST(root="data", train=True, download=True, transform=ToTensor())
test_data = datasets.MNIST(root="data", train=False, download=True, transform=ToTensor())

其中，root参数表示数据集存储的文件夹路径；train参数表示是否下载训练集；download参数表示是否下载数据集；transform参数表示对数据集进行的预处理操作。

2.2、手动下载MNIST数据集

若PyTorch自带的函数下载速度较慢或无法下载，可以手动下载MNIST数据集。下载链接为：

http://yann.lecun.com/exdb/mnist/

下载后的数据集包含以下四个文件：

train-images-idx3-ubyte.gz：训练图像数据集文件。

train-labels-idx1-ubyte.gz：训练标签数据集文件。

t10k-images-idx3-ubyte.gz：测试图像数据集文件。

t10k-labels-idx1-ubyte.gz：测试标签数据集文件。

下载完成后，将数据集文件解压到对应的文件夹中，并使用以下代码进行读取和处理。

import torch

import gzip
import numpy as np
from torchvision.transforms import ToTensor
# 读取.gz压缩文件
def read_image_file(file_path):
    with gzip.open(file_path, 'rb') as f:
        buffer = f.read()
        buffer = np.frombuffer(buffer, dtype=np.uint8, offset=16)
    return buffer.reshape(-1, 28, 28)
def read_label_file(file_path):
    with gzip.open(file_path, 'rb') as f:
        buffer = f.read()
        buffer = np.frombuffer(buffer, dtype=np.uint8, offset=8)
    return buffer
# 根据文件路径读取图像和标签文件
def read_data_images(file_paths):
    images = []
    for file_path in file_paths:
        images.append(read_image_file(file_path))
    return np.concatenate(images, axis=0)
def read_data_labels(file_paths):
    labels = []
    for file_path in file_paths:
        labels.append(read_label_file(file_path))
    return np.concatenate(labels, axis=0)
# 创建MNIST数据集
train_images = read_data_images(['data/train-images-idx3-ubyte.gz'])
train_labels = read_data_labels(['data/train-labels-idx1-ubyte.gz'])
test_images = read_data_images(['data/t10k-images-idx3-ubyte.gz'])
test_labels = read_data_labels(['data/t10k-labels-idx1-ubyte.gz'])
# 将数据转换为Tensor类型并进行预处理
train_data = list(zip(ToTensor()(train_images), torch.LongTensor(train_labels)))
test_data = list(zip(ToTensor()(test_images), torch.LongTensor(test_labels)))

其中read_image_file()和read_label_file()函数用于读取.gz压缩的数据文件，并将数据转换为Numpy数组。read_data_images()和read_data_labels()函数用于读取多个文件的数据，并将它们拼接为一个大数组。

3、创建数据加载器

在使用MNIST数据集时，可以创建数据加载器来加载图像和标签数据，以便于训练和测试模型。

import torch.utils.data as data

# 创建训练数据加载器
train_loader = data.DataLoader(train_data, batch_size=100, shuffle=True)
# 创建测试数据加载器
test_loader = data.DataLoader(test_data, batch_size=100, shuffle=True)

其中，batch_size参数表示每个batch包含的图像数量；shuffle参数表示在每个epoch开始前是否对数据进行随机打乱。

4、构建神经网络模型

在本例中，我们使用全连接神经网络模型，用于对28x28像素的图像进行分类，输出图像所代表的数字值。

import torch.nn as nn

class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.linear1 = nn.Linear(28 * 28, 512)
        self.linear2 = nn.Linear(512, 256)
        self.linear3 = nn.Linear(256, 128)
        self.linear4 = nn.Linear(128, 64)
        self.linear5 = nn.Linear(64, 10)
    def forward(self, x):
        x = x.view(-1, 28 * 28)
        x = nn.functional.relu(self.linear1(x))
        x = nn.functional.relu(self.linear2(x))
        x = nn.functional.relu(self.linear3(x))
        x = nn.functional.relu(self.linear4(x))
        x = self.linear5(x)
        return x
# 实例化神经网络
model = Net()
print(model)

在构建模型时，使用了5个Linear层，其中前4个Linear层的激活函数为ReLU，最后一个Linear层没有激活函数。在forward()函数中，我们使用了view()函数将输入的图像数据展开成1维向量，并依次经过了5个Linear层，最终输出10个数字代表图像的分类结果。

5、定义损失函数和优化器

在训练过程中，需要定义损失函数来计算模型预测值与真实值之间的差距，并通过优化器来更新模型参数，以使损失函数最小化。

import torch.optim as optim

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

在本例中，我们使用交叉熵损失函数和随机梯度下降（SGD）优化器。

6、训练模型

在完成上述步骤后，我们即可开始训练模型。在训练过程中，我们会先使用前向传播得到模型的输出值，再与真实值计算损失函数。然后通过反向传播求得模型对每个参数的梯度，最后使用优化器更新模型的参数。这一过程重复多次，直到模型收敛。

from tqdm import tqdm

# 训练模型
for epoch in range(10):
    total_loss = 0.0
    for images, labels in tqdm(train_loader):
        optimizer.zero_grad()
        outputs = model(images)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()
        total_loss += loss.item() * images.size(0)
    print("Epoch: {}, Train Loss: {:.6f}".format(epoch+1, total_loss/len(train_data)))
print("Finished Training!")

在本例中，我们共进行了10个epoch的训练。在每个epoch中，我们使用tqdm库实现了对训练过程的进度条显示，并统计训练过程中的总损失值。最后输出每个epoch的训练损失值和训练完成的提示信息。

7、测试模型

在训练完成后，我们可使用测试集对模型进行测试，并计算分类准确率。

# 测试模型

correct = 0
total = 0
with torch.no_grad():
    for images, labels in tqdm(test_loader):
        outputs = model(images)
        _, predicted = torch.max(outputs.data, 1)
        total += labels.size(0)
        correct += (predicted == labels).sum().item()
print("Accuracy on test set: {:.2f}%".format(100 * correct / total))

在本例中，我们使用了测试数据集进行了模型测试，并计算了模型的分类准确率。在测试中，我们使用了torch.no_grad()函数避免梯度计算的开销，并使用精度计算公式计算分类准确率。

8、总结

本文通过手把手地演示，讲述了如何在PyTorch中实现MNIST数据集的分类任务。通过读取并预处理数据、构建神经网络模型、定义损失函数和优化器、训练和测试模型，我们可以很方便地完成MNIST数据集的分类，并得出分类准确率。