使用简单MLP实现0-9数字识别，数据集为MNIST

news2024/10/7 5:28:58

简介

需求：
1. 基于pytorch实现简单MLP，完成数字识别，
2. 采用MNIST手写数字作为数据集，MNIST：有6万张训练图片，1万张测试图片
3. 训练结束后，随机取3张测试图片，展示模型的预测结果和真实图片
模型结构：
1. 一层线性层作为输入层，转换输入
2. 中间三层线性层
3. 一层softmax作为输出层，输出结果概率

模型结果

输出结果

Accuracy: 0.0958
epoch:  0 Accuracy:  0.9406
epoch:  1 Accuracy:  0.9611
epoch:  2 Accuracy:  0.9631
epoch:  3 Accuracy:  0.9687
epoch:  4 Accuracy:  0.9703
epoch:  5 Accuracy:  0.9691
epoch:  6 Accuracy:  0.9737
epoch:  7 Accuracy:  0.974
epoch:  8 Accuracy:  0.9717
epoch:  9 Accuracy:  0.9731
Prediction:  tensor(7)
Prediction:  tensor(5)
Prediction:  tensor(4)

真实的测试图片

代码及注释

import torch
from torch.utils.data import DataLoader
from torchvision import transforms
from torchvision.datasets import MNIST
import matplotlib.pyplot as plt

# 定义模型结构
class Net(torch.nn.Module):
    # 初始化模型
    def __init__(self):
        super().__init__()
        self.conv1 = torch.nn.Linear(28*28, 64) # 输入维度为28*28，输出维度为64，第一个全连接层
        self.conv2 = torch.nn.Linear(64, 64) # 输入维度为64，输出维度为64，第二个全连接层
        self.conv3 = torch.nn.Linear(64, 64) # 输入维度为64，输出维度为64，第三个全连接层
        self.conv4 = torch.nn.Linear(64, 10) # 输入维度为64，输出维度为10，第四个全连接层，最后输出10个类别

    def forward(self, x):
        x = torch.relu(self.conv1(x))
        x = torch.relu(self.conv2(x))
        x = torch.relu(self.conv3(x))
        # 由于我们的模型是一个多分类问题，所以我们需要使用softmax来归一化输出最后的概率，加上log是为了防止数值过小，导致数值溢出
        # 由于输出的形状为(batch_size,P)，dim=1表示对第一个维度进行softmax操作，即对概率P进行softmax操作
        x = torch.nn.functional.log_softmax(self.conv4(x), dim=1)
        return x
# 定义加载数据函数
def get_data_loader(is_train):
    dataset = MNIST(root='./data', train=is_train, download=True, transform=transforms.ToTensor())
    # shuffle打乱原数据集，并按64个为一批次，返回到data_loader中
    # 数据加载器会按照指定的批次大小，从数据集中逐批加载数据。
    # 所以在每个迭代中，数据加载器会提供一个批次的数据，但是数据加载器会动态地在每次迭代中加载下一个批次的数据，直到遍历完整个数据集。
    data_loader = DataLoader(dataset, batch_size=64, shuffle=True)
    # 返回数据加载器/迭代器
    return data_loader
# 定义测试函数
def evaluate(test_dataloader, net):
    correct = 0
    total = 0
    # 由于在测试部分，所以取消梯度更新
    with torch.no_grad():
        # 由于我们是在每个data_loader中，对模型进行参数更新，而模型是每个迭代进行一次更新，所以我们将data_loader叫做迭代器
        # 从迭代器，中获取一个批次的数据，返回给x,y
        for x,y in test_dataloader:
            # 由于我们的模型结构的第一层为linear(28*28,64)，而数据集中的样本形状为(batch_size,1,28,28)，1表示通道数，灰度图像只有一个通道
            # 所以我们需要将输入x的形状转换为(batch_size,28*28)，-1表示自动计算，可以使用reshape或者view来进行形状转换
            # 随后将x输入到模型中，得到输出outputs
            outputs = net(x.reshape(-1, 28*28))
            # enumerate(outputs) 是一个 Python 内置函数，它用于将一个可迭代对象（如列表、元组、字符串等）包装成一个枚举对象，同时返回一个索引和对应的值。
            # 例如，outputs = [0.9, 0.8, 0.7]，那么 enumerate(outputs) 的结果是 [(0, 0.9), (1, 0.8), (2, 0.7)]
            for i,output in enumerate(outputs):
                # torch.argmax(output) 返回output中最大值的索引，由于我们的输出维度为10，所以刚好为预测的类别(0-9)
                # 如果相同，那么预测正确，correct+1，同时不论正确与否，预测总数total+1
                if torch.argmax(output) == y[i]:
                    correct += 1
                total += 1
    return correct/total

if __name__ == '__main__':
    # 创建训练和测试数据加载器
    train_dataloader = get_data_loader(is_train=True)
    test_dataloader = get_data_loader(is_train=False)
    # 初始化神经网络
    net = Net()
    # 输出最开始时模型的准确率
    print("Accuracy:", evaluate(test_dataloader, net))
    # 定义优化器
    optimizer = torch.optim.Adam(net.parameters(), lr=0.001)

    # 开始训练
    for epoch in range(10):
        for x,y in train_dataloader:
            # 重置梯度
            optimizer.zero_grad()
            # 正向传播
            outputs = net(x.reshape(-1, 28*28))
            # 计算损失
            # 定义损失函数
            # 由于我们的模型输出为log_softmax，所以我们使用负对数似然损失函数
            # 同时，由于nll_loss不支持无参数传入，所以我们在此处定义损失函数，并同时传入outputs和y
            loss = torch.nn.functional.nll_loss(outputs, y)
            # 反向传播
            loss.backward()
            # 更新参数
            optimizer.step()
        # 每个epoch完成后，都输出模型的准确率
        print('epoch: ',epoch,"Accuracy: ", evaluate(test_dataloader, net))
    # 训练完成后，随机选取三张图像进行预测
    for images, labels in test_dataloader:
        for i in range(3):
            output = net(images[i].reshape(-1, 28*28))
            print("Prediction: ", torch.argmax(output))
            plt.imshow(images[i].view(28, 28), cmap='gray')
            plt.show()
        # print("Image Shape:", images.shape)  # 打印图像的形状
        # # print("Label:", labels.item())  # 打印标签
        # # print("Tensor Data:", images)  # 打印图像的张量数据
        # # plt.imshow(images.view(28, 28), cmap='gray')  # 显示图像
        # # plt.show()
        # outputs = net(images.reshape(-1, 28 * 28))
        # print(outputs.shape)
        # print(outputs)
        break