pytorch 一个完整的训练

训练train.py

test.py按照验证的写法即可
模型使用的是之前的cifar10_seq模型

步骤如下:

  1. 设置使用的设备(cuda、mps)
  2. 创建数据集(数据集自定义参照dataset一文)
  3. 加载数据集
  4. 创建网络并放到设备上、损失函数、优化器以及其相关参数、冻结权重等其他操作(如tensorboard)
  5. 设定整个迭代次数,训练,验证

训练过程:

  1. 设置训练模式
  2. 得到数据与标签,放到要使用的设备
  3. 得到预测结果,计算损失
  4. 清空优化器梯度,反向传播得到梯度,梯度更新
  5. 一些训练进度、时长、损失、精度等自定义显示

验证过程设置验证模式,不计算梯度,不用对优化器进行操作,其他与训练过程一样

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()





你可能感兴趣的:(pytorch,pytorch,深度学习,python)