PyTorch MNIST 手写数字识别教程 - 从入门到精通 - 常规

import torch
from torchvision.datasets import MNIST
from torchvision.transforms import Compose, ToTensor, Normalize
from torch.utils.data import DataLoader
import torch.nn as nn
import torch.nn.functional as F
from torch.optim import Adam
import os
import torch
import numpy as np

BATCH_SIZE = 128
TEST_BATCH_SIZE = 1000


# 准备数据集
def get_dataloader(train, batch_size=BATCH_SIZE):
    transform_fn = Compose([
        ToTensor(),
        Normalize(mean=(0.1307,), std=(0.3081,))
    ])  # mean和std的形状与通道数相同

    dataset = MNIST(root='./data', train=train, transform=transform_fn)
    data_loader = DataLoader(dataset, batch_size=batch_size, shuffle=True)
    return data_loader


class MnistModel(nn.Module):
    def __init__(self):
        super(MnistModel, self).__init__()  # 继承
        self.fc1 = nn.Linear(1 * 28 * 28, 28)  # 参数是input和output的feature
        self.fc2 = nn.Linear(28, 10)

    def forward(self, input):
        # 1.进行形状的修改
        x = input.view([-1, 1 * 28 * 28])  # -1表示根据形状自动调整，也可以改为input.size(0)
        # 2.进行全连接的操作
        x = self.fc1(x)
        # 3.激活函数的处理
        x = F.relu(x)  # 形状没有变化
        # 4.输出层
        out = self.fc2(x)
        return F.log_softmax(out, dim=-1)


model = MnistModel()
optimizer = Adam(model.parameters(), lr=0.001)


def train(epoch):  # epoch表示几轮
    data_loader = get_dataloader(True)  # 获取数据加载器
    for idx, (input, target) in enumerate(data_loader):  # idx表示data_loader中的第几个数据，元组是data_loader的数据
        optimizer.zero_grad()  # 将梯度置0
        output = model(input)  # 调用模型，得到预测值
        loss = F.nll_loss(output, target)  # 调用损失函数，得到损失,是一个tensor
        loss.backward()  # 反向传播
        optimizer.step()  # 梯度的更新
        if idx % 10 == 0:
            print('Epoch:', epoch, 'Iteration:', idx, 'Loss:', loss.item())
    # for是每一轮中的数据进行遍历


def test():
    loss_list = []
    acc_list = []
    test_dataloader = get_dataloader(train=False, batch_size=TEST_BATCH_SIZE)  # 获取测试集
    for idx, (input, target) in enumerate(test_dataloader):
        with torch.no_grad():  # 不计算梯度
            output = model(input)
            cur_loss = F.nll_loss(output, target)
            loss_list.append(cur_loss)
            # 计算准确率，output大小[batch_size,10] target[batch_size] batch_size是多少组数据，10列是每个数字概率
            pred = output.max(dim=-1)[-1]  # 获取最大值位置
            cur_acc = pred.eq(target).float().mean()
            acc_list.append(cur_acc)
    print('平均准确率：', np.mean(acc_list), '平均损失：', np.mean(loss_list))

if __name__ == '__main__':
    test()
    for i in range(3):  # 训练三轮
        train(i)
    test()

代码说明：

导入必要的库: 包括 PyTorch、Torchvision、NumPy 等库。
定义超参数: 包括批量大小、测试批量大小等参数。
准备数据集: 使用 MNIST 数据集，并进行数据增强和标准化。
构建模型: 使用两个全连接层构建简单的 MNIST 模型。
定义优化器: 使用 Adam 优化器来优化模型参数。
训练函数: 循环训练模型，并记录训练过程中的损失等信息。
测试函数: 使用测试集评估模型的性能，并记录准确率和损失。
运行程序: 首先测试模型，然后训练模型三轮，最后再次测试模型。

运行方法：

确保已安装 PyTorch 和 Torchvision 库。
在代码所在的目录下创建一个名为 “data” 的文件夹。
运行代码，程序会自动下载 MNIST 数据集到 “data” 文件夹中。
等待训练完成后，程序会输出模型的最终测试性能。

其他优化建议：

可以尝试使用更复杂的模型架构，例如卷积神经网络。
可以调整超参数，例如学习率、批量大小等，以提高模型性能。
可以使用 GPU 加速训练过程，提高训练速度。
可以使用 TensorBoard 等工具可视化训练过程，帮助理解模型的训练过程。

希望这篇文章能帮助您入门 PyTorch MNIST 手写数字识别！