Pytorch实现MNIST手写数字识别

造曲率引擎

于 2024-06-08 16:52:36 发布

阅读量870

点赞数 17

分类专栏：机器学习文章标签： pytorch 人工智能 python

本文链接：https://blog.csdn.net/qq_40469664/article/details/138970996

版权

机器学习专栏收录该内容

3 篇文章 0 订阅

订阅专栏

Pytorch实现MNIST手写数字识别

《机器学习》课程布置了一个简单的作业任务，需要基于Pytorch实现MNIST手写数字识别。借此机会梳理下Pytorch实现的细节。

环境和调用

本项目在windows11平台运行，语言Python3.9，pytorch版本2.2.1+cu121。
Python调用：

# mnist_train.py
import numpy as np
import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import datasets, transforms
from torch.utils.data import DataLoader, Subset
import matplotlib.pyplot as plt

# model.py
import torch
import torch.nn as nn
import torch.nn.functional as F

加载数据集

MNIST 数据集是经典手写数字图片数据集，官方下载地址为MNIST。不过有时可能没法正常下载。
在PyTorch中，可以使用 torchvision 库轻松下载和加载 MNIST 数据集。

torchvision 是 PyTorch 生态系统中的一个重要库，专门用于计算机视觉任务。它提供了一系列工具和功能，简化了加载、预处理和增强图像数据的过程，同时也提供了预训练的深度学习模型，方便开发者进行迁移学习和模型验证。

class CustomNormalize(object):
    """
    自定义转换操作，将像素值从 [0, 1] 范围转换到 [-0.5, 0.5] 范围
    """

    def __call__(self, tensor):
        return tensor - 0.5

# 构建pipeline，转换为 PyTorch 张量, 像素值归一化到 [-0.5, 0.5] 范围
pipeline = transforms.Compose([
    transforms.ToTensor(),     
    CustomNormalize()
])

# 下载数据集
full_train_dataset = datasets.MNIST("./data",train=True,download=True,transform=pipeline)
test_dataset = datasets.MNIST("./data",train=False,download=True,transform=pipeline)

# 从原训练集每个类别随机选取300个样本作为训练集
targets = np.array(full_train_dataset.targets)
class_indices = {i: np.where(targets == i)[0] for i in range(10)}
selected_indices = []
for i in range(10):
    selected_indices.extend(np.random.choice(class_indices[i], 300, replace=False))
train_dataset = Subset(full_train_dataset, selected_indices)

# 加载数据
train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
test_loader = DataLoader(test_dataset, batch_size=1000, shuffle=False)
print(f'训练样本数量: {len(train_loader.dataset)}')

torchvision.datasets

torchvision.datasets 可用于下载数据集。

full_train_dataset = datasets.MNIST("./data",train=True,download=True,transform=pipeline)
test_dataset = datasets.MNIST("./data",train=False,download=True,transform=pipeline)

这段代码将MNIST数据集下载到文件夹"./data"中（如果文件夹不存在，会自动创建）。

train=区分下载的是训练集、测试集
transform=定义数据转换处理操作，需要传递一个用于预处理和增强数据的转换管道。这个转换管道通常是组合由 torchvision.transforms.Compose 提供的多个转换操作实现。

torchvision.transforms

torchvision.transforms 可对数据集进行转换优化。
它提供了不少转换操作，常用的例如：

transforms.ToTensor()：将 PIL 图像或 NumPy 数组转换为 PyTorch 张量，并且将像素值从 [0, 255] 范围缩放到 [0, 1] 范围。
transforms.Normalize(mean, std)：对图像的每个通道进行标准化，使其具有指定的均值和标准差。这通常用于加速模型的收敛过程。

在本项目中，要求将像素值归一化到 [-0.5, 0.5] 范围。因此先调用transforms.ToTensor()，将图像转换为 PyTorch 张量，将像素值从 [0, 255] 范围缩放到 [0, 1] 范围。然后调用自定义的转换CustomNormalize()，将 [0, 1] 范围转换到 [-0.5, 0.5] 范围。

class CustomNormalize(object):
    def __call__(self, tensor):
        return tensor - 0.5

pipeline = transforms.Compose([
    transforms.ToTensor(),     
    CustomNormalize()
])

torch.utils.data.Subset

torch.utils.data.Subset 是 PyTorch 中的一个工具类，用于创建一个给定数据集的子集。
本项目要求在0~9每一类训练样本中各随机选择300个（共3000个）作为实际训练集。

targets = np.array(full_train_dataset.targets)
class_indices = {i: np.where(targets == i)[0] for i in range(10)}
selected_indices = []
for i in range(10):
    selected_indices.extend(np.random.choice(class_indices[i], 300, replace=False))
train_dataset = Subset(full_train_dataset, selected_indices)

torch.utils.data.DataLoader

torch.utils.data.DataLoader 用于加载已经下载的数据集。

train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
test_loader = DataLoader(test_dataset, batch_size=1000, shuffle=False)

batch_size= 设置数据批次大小。
- 对于训练集，一般选择较小的批量大小，有助于在内存使用和计算效率之间取得平衡。较小的批量大小可以更频繁地更新模型参数，从而更快地收敛。
- 对于测试集，可以选择较大的批量大小。因为测试阶段不需要进行反向传播和参数更新，只需要前向传播来计算模型的性能。较大的批量大小可以更高效地利用硬件资源，加快计算速度。
shuffle= 设置数据是否需要被打乱顺序。对于训练集，在每个训练周期（epoch）开始时打乱数据顺序。这是非常重要的，因为打乱数据可以防止模型在训练过程中学习到数据的固定顺序，从而提高模型的泛化能力。数据顺序的随机化有助于避免模型在相邻样本上产生依赖性，并能更好地泛化到未见过的数据。

模型构建

class DNN(nn.Module):
    def __init__(self):
        super(DNN, self).__init__()
        self.conv1 = nn.Conv2d(in_channels=1, out_channels=16, kernel_size=5, padding=0, stride=1)
        self.pool = nn.MaxPool2d(kernel_size=2, stride=2)
        self.conv2 = nn.Conv2d(in_channels=16, out_channels=32, kernel_size=5, padding=0, stride=1)
        self.fc1 = nn.Conv2d(in_channels=32, out_channels=100, kernel_size=4, padding=0, stride=1)
        self.fc2 = nn.Linear(100, 100)
        self.fc3 = nn.Linear(100, 10)

    def forward(self, x):
        x = F.relu(self.conv1(x))
        x = self.pool(x)
        x = F.relu(self.conv2(x))
        x = self.pool(x)
        x = F.relu(self.fc1(x))
        x = x.view(-1, 100)
        x = F.relu(self.fc2(x))
        x = self.fc3(x)
        return F.softmax(x, dim=1)

模型构建要根据实际情况选择，这里不过多赘述。
简单复习下卷积层输出尺寸的计算：
$输出 = [(输入 + 2 * 边缘填充 - 卷积核大小) / 步幅] + 1$

此外，注意 Pytorch 中卷积层 torch.nn.Conv2d 得到的张量尺寸为 [batch_size, channels，height，width]，而全连接层 torch.nn.Linear 的输入应该是二维张量。如果要将卷积层的输出作为全连接层的输入，需要通过 view() 函数调整尺寸。

很多时候我们会习惯直接 copy 网上现成的 model，但为了能够有独立构建一个模型的能力，还是需要充分理解神经网络中的数学基础。

torch.nn.Module

torch.nn.Module 是 PyTorch 中神经网络模块的基类。它提供了一种方便的方法来组织神经网络模型，并管理模型参数、前向传播以及一些其他功能。

在Python中，super() 函数用于调用父类的方法。在 torch.nn.Module 的子类中，通常会在子类的构造函数 __init__ 中调用父类的构造函数，以确保父类中定义的一些初始化逻辑也被执行。这通常通过 super() 函数来实现。

super(子类名, self).__init__()

子类名: 指定子类的名称。
self: 表示当前对象的引用，用于传递当前实例给父类构造函数。
__init__(): 父类的构造函数。

简单模型构建示例：

class SimpleModel(nn.Module):
    def __init__(self):
        super(SimpleModel, self).__init__()
        self.fc1 = nn.Linear(784, 128)
        self.relu = nn.ReLU()
        self.fc2 = nn.Linear(128, 10)

    def forward(self, x):
        x = torch.flatten(x, 1)
        x = self.fc1(x)
        x = self.relu(x)
        x = self.fc2(x)
        return x

模型训练

# 创建模型、损失函数和优化器
model = DNN()
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=learning_rate)

# 用于记录损失
train_losses = []

# 训练模型
for epoch in range(num_epochs):
    model.train()
    running_loss = 0.0
    for images, labels in train_loader:
        optimizer.zero_grad()
        outputs = model(images)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()
        running_loss += loss.item()
    loss = running_loss / len(train_loader)
    train_losses.append(loss)
    print(f'Epoch [{epoch + 1}/{num_epochs}], Loss: {loss:.4f}')

训练细节

for epoch in range(num_epochs):
    model.train()
    for images, labels in train_loader:
        optimizer.zero_grad()
        outputs = model(images)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()

model.train()：这是 PyTorch 中的一个方法，用于设置模型为训练模式。在训练模式下，模型的行为可能会与评估模式下有所不同，例如启用了 Batch Normalization 和 Dropout。
optimizer.zero_grad()：在每次迭代开始时，用于将模型参数的梯度归零。这是因为 PyTorch 默认会累积梯度，而在每次迭代时我们需要重新计算梯度。
loss.backward()：执行反向传播算法，计算模型参数的梯度。
optimizer.step()：基于计算得到的梯度更新模型参数。

完整代码

mnist_train.py

"""
实现MNIST分类训练测试
"""
import numpy as np
import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import datasets, transforms
from torch.utils.data import DataLoader, Subset
import matplotlib.pyplot as plt

from model import DNN


# 固定随机数种子
seed = 42
np.random.seed(seed)
torch.manual_seed(seed)

# 超参
batch_size = 64
learning_rate = 0.001
num_epochs = 20

class CustomNormalize(object):
    """
    自定义转换操作，将像素值从 [0, 1] 范围转换到 [-0.5, 0.5] 范围
    """

    def __call__(self, tensor):
        return tensor - 0.5


def drawCurve(losses, tag):
    """
    绘制loss的拟合曲线
    """
    epochs = np.arange(1, len(losses) + 1)
    # 多项式拟合
    coefficients = np.polyfit(epochs, losses, 10)
    polynomial = np.poly1d(coefficients)
    fit_line = polynomial(epochs)

    # 绘制训练损失曲线和拟合曲线
    plt.figure(figsize=(10, 5))
    plt.plot(epochs, train_losses, label=f'{tag} Loss', marker='o')
    plt.plot(epochs, fit_line, label='Fitted Curve', linestyle='--')
    plt.xlabel('Epoch')
    plt.ylabel('Loss')
    plt.title(f'the curve of the {tag} loss')
    plt.xticks(epochs)
    plt.legend()
    plt.show()


# 构建pipeline，转换为 PyTorch 张量, 像素值归一化到 [-0.5, 0.5] 范围
pipeline = transforms.Compose([
    transforms.ToTensor(),     
    CustomNormalize()
])

# 下载数据集
full_train_dataset = datasets.MNIST("./data",train=True,download=True,transform=pipeline)
test_dataset = datasets.MNIST("./data",train=False,download=True,transform=pipeline)

# 从原训练集每个类别随机选取300个样本作为训练集
targets = np.array(full_train_dataset.targets)
class_indices = {i: np.where(targets == i)[0] for i in range(10)}
selected_indices = []
for i in range(10):
    selected_indices.extend(np.random.choice(class_indices[i], 300, replace=False))
train_dataset = Subset(full_train_dataset, selected_indices)

# 加载数据
train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
test_loader = DataLoader(test_dataset, batch_size=1000, shuffle=False)
print(f'训练样本数量: {len(train_loader.dataset)}')

# 创建模型、损失函数和优化器
model = DNN()
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=learning_rate)

# 用于记录损失
train_losses = []

# 训练模型
for epoch in range(num_epochs):
    model.train()
    running_loss = 0.0
    for images, labels in train_loader:
        optimizer.zero_grad()
        outputs = model(images)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()
        running_loss += loss.item()
    loss = running_loss / len(train_loader)
    train_losses.append(loss)
    print(f'Epoch [{epoch + 1}/{num_epochs}], Loss: {loss:.4f}')

# 测试模型
model.eval()
correct = 0
total = 0
test_loss = 0.0
with torch.no_grad():
    for images, labels in test_loader:
        outputs = model(images)
        loss = criterion(outputs, labels)
        test_loss += loss.item()
        _, predicted = torch.max(outputs.data, 1)
        total += labels.size(0)
        correct += (predicted == labels).sum().item()

print(f'测试集上的loss: {test_loss / len(test_loader)}, 准确率: {100 * correct / total:.2f}%')

# 绘制训练和测试的损失曲线
drawCurve(train_losses, 'training')

model.py

"""
DNN模型构建
"""
import torch
import torch.nn as nn
import torch.nn.functional as F


class DNN(nn.Module):
    def __init__(self):
        super(DNN, self).__init__()
        self.conv1 = nn.Conv2d(in_channels=1, out_channels=16, kernel_size=5, padding=0, stride=1)
        self.pool = nn.MaxPool2d(kernel_size=2, stride=2)
        self.conv2 = nn.Conv2d(in_channels=16, out_channels=32, kernel_size=5, padding=0, stride=1)
        self.fc1 = nn.Conv2d(in_channels=32, out_channels=100, kernel_size=4, padding=0, stride=1)
        self.fc2 = nn.Linear(100, 100)
        self.fc3 = nn.Linear(100, 10)

    def forward(self, x):
        x = F.relu(self.conv1(x))
        x = self.pool(x)
        x = F.relu(self.conv2(x))
        x = self.pool(x)
        x = F.relu(self.fc1(x))
        x = x.view(-1, 100)
        x = F.relu(self.fc2(x))
        x = self.fc3(x)
        return F.softmax(x, dim=1)