神经网络入门实战:(二十)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("模型已保存!!")

结果如下:


上一篇 下一篇
神经网络入门实战(十九) 待发布
相关推荐
IT_陈寒10 小时前
React 18实战:7个被低估的Hooks技巧让你的开发效率提升50%
前端·人工智能·后端
逛逛GitHub11 小时前
飞书多维表“独立”了!功能强大的超出想象。
人工智能·github·产品
机器之心11 小时前
刚刚,DeepSeek-R1论文登上Nature封面,通讯作者梁文锋
人工智能·openai
aneasystone本尊14 小时前
学习 Chat2Graph 的知识库服务
人工智能
IT_陈寒14 小时前
Redis 性能翻倍的 7 个冷门技巧,第 5 个大多数人都不知道!
前端·人工智能·后端
飞哥数智坊1 天前
GPT-5-Codex 发布,Codex 正在取代 Claude
人工智能·ai编程
倔强青铜三1 天前
苦练Python第46天:文件写入与上下文管理器
人工智能·python·面试
虫无涯1 天前
Dify Agent + AntV 实战:从 0 到 1 打造数据可视化解决方案
人工智能
Dm_dotnet1 天前
公益站Agent Router注册送200刀额度竟然是真的
人工智能
算家计算1 天前
7B参数拿下30个世界第一!Hunyuan-MT-7B本地部署教程:腾讯混元开源业界首个翻译集成模型
人工智能·开源