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