训练数据每张图片是 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/
您可以选择一种方式赞助本站
支付宝扫一扫赞助
微信钱包扫描赞助
赏