训练数据每张图片是 28*28的，所以输入维度为 784，输出维度为10 (即表示0-9的数字)

可以通过两个隐藏层，以batch_size=1为例

输入 1*784 变成 1*1000，中间需要乘以一个 784*1000 的矩阵，体现到代码中就是 nn.Linear(784, 1000)

后面再从1000个通道转到500，再到10

一、数据下载

https://www.kaggle.com/datasets/shreyasvaderiyattils/mnist-npy-dataset?resource=download

二、训练模型

import numpy as np
import torch

# 1. 准备数据
train_x = np.load('data/MNIST_train_images.npy', allow_pickle=True)
train_y = np.load('data/MNIST_train_labels.npy', allow_pickle=True)
# print(train_x.shape) # (50000, 784)
# print(train_y.shape) # (50000,)
x_train = torch.from_numpy(train_x).float()
y_train = torch.from_numpy(train_y).long()

x_train = x_train.cuda()
y_train = y_train.cuda()

# 2. 搭建全连接神经网络
import torch.nn as nn
import torch.nn.functional as F


class MNIST(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc1 = nn.Linear(784, 1000)  # 784 -> 1000
        self.fc2 = nn.Linear(1000, 500)  # 1000 -> 500
        self.fc3 = nn.Linear(500, 10)  # 500 -> 10

    def forward(self, x):
        x = F.relu(self.fc1(x))  # 先进行线性变换，再进行激活函数
        x = F.relu(self.fc2(x))  # 同样，先做线性变换，再做激活函数
        x = self.fc3(x)  # 待会儿会调用交叉熵损失函数，里面会调用softmax函数，所以这里不用调用softmax函数
        return x


# 3. 训练模型
mnist = MNIST().cuda()
for i in range(100):
    y = mnist(x_train)  # 3.1 前向传播
    loss = F.cross_entropy(y, y_train)  # 3.2 计算损失函数
    loss.backward()  # 3.3 反向传播

    lr = 0.001  # 3.4 定义学习率
    with torch.no_grad():  # 3.5 更新权重参数
        for param in mnist.parameters():
            param -= lr * param.grad  # 参数新值 = 参数旧值 - 学习率 * 梯度

    mnist.zero_grad()  # 3.6 梯度清零
    print(i, loss)

# 4. 保存模型
torch.save(mnist, 'model/mnist_model.pkl')

三、调用模型

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F


# 1. 定义或引入神经网络模型(为了简单，这里就不单独写一个文件定义模型了)
class MNIST(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc1 = nn.Linear(784, 1000)  # 784 -> 1000
        self.fc2 = nn.Linear(1000, 500)  # 1000 -> 500
        self.fc3 = nn.Linear(500, 10)  # 500 -> 10

    def forward(self, x):
        x = F.relu(self.fc1(x))  # 先进行线性变换，再进行激活函数
        x = F.relu(self.fc2(x))  # 同样，先做线性变换，再做激活函数
        x = self.fc3(x)  # 待会儿会调用交叉熵损失函数，里面会调用softmax函数，所以这里不用调用softmax函数
        return x


# 2. 加载模型
mnist = torch.load("model/mnist_model.pkl")
mnist = mnist.cuda()

# 3. 验证模型
# 准备数据
test_x = np.load('data/MNIST_test_images.npy', allow_pickle=True)
test_y = np.load('data/MNIST_test_labels.npy', allow_pickle=True)
x_test = torch.from_numpy(test_x).float()
y_test = torch.from_numpy(test_y).long()

# 验证10个样本
for i in range(10):
    y = mnist(x_test[i].cuda())
    print(y)
    label = y_test[i]
    predict = torch.argmax(y)
    print('答案：', label)
    print('预测：', predict)
    print(predict == label if '预测正确' else '预测错误')
    print('-------------------')

参考资料：https://www.bilibili.com/video/BV17w411k7a9/