训练train.py
test.py按照验证的写法即可
模型使用的是之前的cifar10_seq模型
步骤如下:
- 设置使用的设备(cuda、mps)
- 创建数据集(数据集自定义参照dataset一文)
- 加载数据集
- 创建网络并放到设备上、损失函数、优化器以及其相关参数、冻结权重等其他操作(如tensorboard)
- 设定整个迭代次数,训练,验证
训练过程:
- 设置训练模式
- 得到数据与标签,放到要使用的设备
- 得到预测结果,计算损失
- 清空优化器梯度,反向传播得到梯度,梯度更新
- 一些训练进度、时长、损失、精度等自定义显示
验证过程设置验证模式,不计算梯度,不用对优化器进行操作,其他与训练过程一样
import torch
from torch import nn
import torchvision
from torch.utils.tensorboard import SummaryWriter
from torch.utils.data import DataLoader
import time
from P9_model import *
# 关于gpu,网络模型、数据需要从CPU挪到GPU训练
# 数据需要挪到GPU去计算,模型中的参数也是需要GPU计算。
# 与之不同的是损失函数虽然也在nn.Module中,但他没有parameter或buffer,因此不需要放到GPU
device = torch.device('mps' if torch.has_mps else 'cpu')
# device = torch.device('cpu')
# device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
print('训练在',device,'上进行')
# 数据集
train_dataset = torchvision.datasets.CIFAR10('./data', train=True, transform=torchvision.transforms.ToTensor(),
download=True)
val_dataset = torchvision.datasets.CIFAR10('./data', train=False, transform=torchvision.transforms.ToTensor(),
download=True)
train_data_size = len(train_dataset)
val_data_size = len(val_dataset)
print('训练集数据长度为%d' % train_data_size)
print('测试集数据长度为%d' % val_data_size)
# 加载数据集
train_dataloader = DataLoader(train_dataset, batch_size=64)
val_dataloader = DataLoader(val_dataset, batch_size=64)
# 创建网络,这里使用之前的模型,一般情况下模型也是单写在其他文件中,需要import
model = cifar10_seq()
model.to(device)
# 损失函数
loss_fn = nn.CrossEntropyLoss()
# 优化器
lr = 1e-2
optimizer = torch.optim.SGD(model.parameters(), lr)
# 训练与测试次数
total_train_step = 0
total_val_step = 0
# 迭代次数
epoch = 10
# tensorboard
writer = SummaryWriter('./p9')
for i in range(epoch):
print('————————————第%d轮训练开始————————————' % i)
start = time.time()
# 训练
# 当模型有dropout和BN等,调成训练模式是必须的
model.train()
total_train_loss = 0
for data in train_dataloader:
imgs, targets = data
# 数据放到GPU需要更新变量
imgs = imgs.to(device)
targets = targets.to(device)
outputs = model(imgs)
# loss的结果是tensor类型,loss.item()的结果是float类型
loss = loss_fn(outputs, targets)
optimizer.zero_grad()
loss.backward()
optimizer.step()
total_train_loss += loss
writer.add_scalar('train_loss', total_train_loss.item(), i)
print('训练loss: %f' % total_train_loss.item())
# 也可以输出每个batch的loss,或每100batch输出一次loss,都可以自定义
end_train = time.time()
print('训练时长: %f' % (end_train-start))
# 验证
# 当模型有dropout和BN等,调成验证模式是必须的
model.eval()
total_val_loss = 0
total_acc = 0
with torch.no_grad():
for data in val_dataloader:
imgs, targets = data
imgs = imgs.to(device)
targets = targets.to(device)
outputs = model(imgs)
loss = loss_fn(outputs, targets)
total_val_loss += loss
# 对于分类任务可以计算正确率,即预测类别等于真正类别的数量
accuracy = (outputs.argmax(1) == targets).sum()
total_acc += accuracy
writer.add_scalar('val_loss', total_val_loss.item(), i)
print('验证loss: %f' % total_val_loss.item())
writer.add_scalar('val_accuracy', total_acc/val_data_size, i)
print('测试集正确率:%f' % (total_acc / val_data_size))
end_val = time.time()
print('验证时长: %f' % (end_val - end_train))
# # 可以每周期保存训练的模型,也可以保存间隔自定义
# torch.save(model, 'model%d' % i)
# print('模型已保存')
writer.close()