import torch.optim
import torchvision
from torch import nn
from torch.utils.data import DataLoader
from torch.utils.tensorboard import SummaryWriter
from torchvision.transforms import transforms
from model import *
# 准备数据集
tran_data = torchvision.datasets.CIFAR10('./data', True, transform=transforms.ToTensor(), download=True)
test_data = torchvision.datasets.CIFAR10('./data', False, transform=transforms.ToTensor(), download=True)
# 查看数据集长度
tran_len = len(tran_data)
test_len = len(test_data)
print(f"训练集长度:{tran_len};测试集长度:{test_len}")
# 加载数据集
tran_dataloader = DataLoader(tran_data, batch_size=64)
test_dataloader = DataLoader(test_data, batch_size=64)
# 创建模型
my_model = Model()
# 损失函数
loss_fun = nn.CrossEntropyLoss()
# 优化器
learning_rate = 0.01
optimize = torch.optim.SGD(my_model.parameters(), lr=learning_rate)
# 设置训练网络的一些参数
# 记录训练的次数
total_train_step = 0
total_test_step = 0
# 设置训练次数
epoch = 10
writer = SummaryWriter('./logs_procedure1')
for i in range(epoch):
print("---------------------")
print(f"第{i + 1}轮训练开始")
for data in tran_dataloader:
imgs, targets = data
output = my_model(imgs)
loss = loss_fun(output, targets)
# 优化模型
optimize.zero_grad()
loss.backward()
optimize.step()
total_train_step += 1
if total_train_step % 100 == 0:
print(f"训练次数:{total_train_step},Loss:{loss.item()}")
writer.add_scalar('train_loss', loss.item(), total_train_step)
# 测试数据集开始
with torch.no_grad():
total_loss = 0
for data in test_dataloader:
imgs, targets = data
output = my_model(imgs)
loss = loss_fun(output, targets)
total_loss = total_loss + loss.item()
print(f"第{i + 1}轮训练,测试集上的误差为{total_loss}")
writer.add_scalar('test_loss',total_loss,i+1)
torch.save(my_model,f'my_model_{i+1}.pth')
writer.close()
模型训练基本套路
最新推荐文章于 2024-08-20 00:18:09 发布