神经网络入门实战:(二十)MNIST训练网络(只用线性层和ReLU)

MNIST训练网络(只用线性层和ReLU)

该数据集一共有7万张图片,其中6万张是训练集,1万张是测试集;每张图片都是28×28像素的单通道(黑白)图片

类比 CIFAR10 的训练过程:

python 复制代码
import torch
from torch import nn
import torchvision
from torchvision import transforms
from torch.utils.data import DataLoader
from torch.utils.tensorboard import SummaryWriter
from NN_models import *


# 检查CUDA是否可用,并设置设备为 GPU
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

dataclass_transform = transforms.Compose([
	transforms.ToTensor(),
	transforms.Normalize((0.1307,), (0.3081,))
])

train_dataset = torchvision.datasets.MNIST(root='E:\\4_Data_sets\\MNIST', train=True,transform=dataclass_transform, download=True)
test_dataset = torchvision.datasets.MNIST(root='E:\\4_Data_sets\\MNIST', train=False,transform=dataclass_transform, download=True)

# 训练和测试数据集的长度
train_data_size = len(train_dataset)
test_size = len(test_dataset)
print(train_data_size,test_size)

train_dataloader = DataLoader(dataset=train_dataset,batch_size=64)
test_dataloader = DataLoader(dataset=test_dataset,batch_size=64)

# 创建网络模型
class MNIST_NET(nn.Module):
	def __init__(self):
		super(MNIST_NET, self).__init__()
		self.model = nn.Sequential(
			nn.Flatten(),
			nn.Linear(784, 512),
			nn.ReLU(),  # 添加ReLU激活函数
			nn.Linear(512, 256),
			nn.ReLU(),  # 添加ReLU激活函数
			nn.Linear(256, 128),
			nn.ReLU(),  # 添加ReLU激活函数
			nn.Linear(128, 64),
			nn.ReLU(),  # 添加ReLU激活函数
			nn.Linear(64, 10)
		)

	def forward(self, x):
		x = self.model(x)
		return x

MNIST_NET_Instance = MNIST_NET().to(device)

# 定义损失函数
loss = nn.CrossEntropyLoss()
# 定义优化器
learning_rate = 0.01
optimizer = torch.optim.SGD(MNIST_NET_Instance.parameters(), lr=learning_rate, momentum=0.9)

# 开始训练
total_train_step = 0
first_train_step = 0
total_test_step = 0
epoch_sum = 10 # 迭代次数

# 添加tensorboard
writer = SummaryWriter('logs')

for i in range(epoch_sum):
	print("------------第 {} 轮训练开始了------------:".format(i+1))

	# 训练步骤开始
	for data in train_dataloader:
		imgs, labels = data
		imgs, labels = imgs.to(device), labels.to(device)  # 将数据和目标移动到GPU
		outputs = MNIST_NET_Instance(imgs)
		loss_real = loss(outputs, labels) # 这里的损失变量 loss_real,千万别和损失函数 loss 相同,否则会报错!
		optimizer.zero_grad()
		loss_real.backward()
		optimizer.step()

		total_train_step += 1
		# 表示第一轮训练结束,取每一轮的第一个batch_size来看看训练效果
		if total_train_step % 938 == 0:
			first_train_step += 1
			print("训练次数为:{}, loss为:{}".format(total_train_step, loss_real)) # 此训练次数非训练轮次,而是训练到第几个batch_size了
			writer.add_scalar('first_batch_size', loss_real.item(), first_train_step)
		writer.add_scalar('total_batch_size', loss_real.item(), total_train_step)


	# 每训练一轮,就使用测试集看看训练效果
	total_test_loss = 0
	with torch.no_grad(): # 后续测试不计算梯度    
		for data in test_dataloader:
			imgs, labels = data
			imgs, labels = imgs.to(device), labels.to(device)
			outputs = MNIST_NET_Instance(imgs)
			loss_fake = loss(outputs, labels)
			total_test_loss += loss_fake.item()
	print("# # 整体测试集上的LOSS为:{}".format(total_test_loss))

writer.close()

torch.save(MNIST_NET_Instance,"E:\\5_NN_model\\MNIST_NET_train10")
print("模型已保存!!")

结果如下:


上一篇 下一篇
神经网络入门实战(十九) 待发布
相关推荐
「、皓子~7 分钟前
后台管理系统的诞生 - 利用AI 1天完成整个后台管理系统的微服务后端+前端
前端·人工智能·微服务·小程序·go·ai编程·ai写作
笑衬人心。42 分钟前
初学Spring AI 笔记
人工智能·笔记·spring
luofeiju1 小时前
RGB下的色彩变换:用线性代数解构色彩世界
图像处理·人工智能·opencv·线性代数
测试者家园1 小时前
基于DeepSeek和crewAI构建测试用例脚本生成器
人工智能·python·测试用例·智能体·智能化测试·crewai
张较瘦_1 小时前
[论文阅读] 人工智能 + 软件工程 | Call Me Maybe:用图神经网络增强JavaScript调用图构建
论文阅读·人工智能·软件工程
大模型真好玩1 小时前
准确率飙升!Graph RAG如何利用知识图谱提升RAG答案质量(四)——微软GraphRAG代码实战
人工智能·python·mcp
xkroy1 小时前
网络协议概念与应用层
网络
Baihai_IDP1 小时前
vec2text 技术已开源!一定条件下,文本嵌入向量可“近乎完美地”还原
人工智能·面试·llm
江太翁1 小时前
Pytorch torch
人工智能·pytorch·python
筏.k1 小时前
C++ 网络编程(14) asio多线程模型IOThreadPool
网络·c++·架构