【学习笔记】【Pytorch】16、模型训练套路

内容概述

本内容主要是介绍一个完整的模型训练套路，以 CIFAR-10 数据集为例。

模型训练步骤：

• 准备数据：创建 datasets 实例
• 加载数据：创建 DataLoader 实例
• 准备模型：神经网络结构
• 设置损失函数
• 设置优化器
• 开始训练：
  • 从 batch_size 个数据中分别取出个图片数据和标签数据
  • 图片数据输入到神经网络模型里后输出训练结果
  • 将训练结果和标签数据经过损失函数
  • 调用优化器实例的梯度清零API
  • 调用反向传播API
  • 调用优化器实例参数优化API
• 开始测试：（使用每轮训练好、但不进行优化的模型）
• 结果聚合展示

代码实现

import torch.optim.optimizer
import torchvision
from torch import nn, optim
from torch.nn import Conv2d, MaxPool2d, Flatten, Linear
from torch.utils.data import DataLoader
from torch.utils.tensorboard import SummaryWriter

from CIFAR_model import Model


# 1.创建 CIFAR10 数据集的训练和测试实例
train_data = torchvision.datasets.CIFAR10(root="./dataset", train=True,
                                          transform=torchvision.transforms.ToTensor(),
                                          download=True)
test_data = torchvision.datasets.CIFAR10(root="./dataset", train=False,
                                          transform=torchvision.transforms.ToTensor(),
                                          download=True)

# Length 长度
train_data_len = len(train_data)
test_data_len = len(test_data)
print(f"训练数据集的长度：{train_data_len}")
print(f"测试数据集的长度：{test_data_len}")

# 2.利用 DataLoader 加载数据集
train_dataloader = DataLoader(train_data, batch_size=64)
test_dataloader = DataLoader(test_data, batch_size=64)

# 3.搭建神经网络：CIFAR-10模型
model = Model()  # 创建实例

# 4.损失函数
loss_fn = nn.CrossEntropyLoss()

# 5.优化器
learning_rate = 0.01  # 1e-2 = 1 x (10)^(-2) = 0.01
optimizer = torch.optim.SGD(model.parameters(), learning_rate)

# 6.设置训练网络的一些参数
total_train_step = 0  # 记录训练的次数
total_test_step = 0  # 记录测试的次数
epoch = 10  # 训练的轮数

writer = SummaryWriter("./CIFAR_logs")  # 创建实例

# 从 DATaLoader 实例取出 batch_size 个数据 -->
# 从 batch_size 个数据中分别取出个图片数据和标签数据 -->
# 图片数据输入到神经网络模型里后输出训练结果 -->
# 将训练结果和标签数据经过损失函数 -->
# 调用优化器实例的梯度清零API -->
# 调用反向传播API -->
# 调用优化器实例参数优化API
for i in range(epoch):
    print(f"--------第{i}轮训练开始--------")

    # 训练步骤开始
    model.tain()  # 模型进入训练模式（仅针对有Dropout,BatchNorm层的网络结构）
    for data in train_dataloader:
        imgs, targets = data
        outputs = model(imgs)
        result_loss = loss_fn(outputs, targets)  # 计算每个参数对应的损失

        optimizer.zero_grad()  # 每个参数对应的梯度清零
        result_loss.backward()  # 反向传播，计算每个参数对应的梯度
        optimizer.step()  # 每个参数根据上一步得到的梯度进行优化

        total_train_step += 1
        if total_train_step % 100 == 0:
            print(f"训练次数：{total_train_step}，Loss：{result_loss.item()}")
            writer.add_scalar("train_loss", result_loss.item(), total_train_step)

    # 测试步骤开始（每轮训练好、但不进行优化的模型）
    model.eval()  # 模型进入测试模式（仅针对有Dropout,BatchNorm层的网络结构）
    total_test_loss = 0
    total_accuracy = 0
    with torch.no_grad():  # 不进行优化的
        for data in test_dataloader:
            imgs, targets = data
            outputs = model(imgs)
            result_loss = loss_fn(outputs, targets)
            total_test_loss += result_loss  # 损失累加

            # 计算标签正确数(取得在 1 方向，概率最大的索引，即得标签输出值；对比标签输出值与目标值+求和：True=1，False=0)
            accuracy = (outputs.argmax(1) == targets).sum()
            total_accuracy += accuracy

    print(f"整体测试集上的Loss：{total_test_loss}")
    print(f"整体测试集上的正确率：{total_accuracy / test_data_len}")
    writer.add_scalar("test_loss", total_test_loss.item(), total_test_step)
    total_test_step += 1

    torch.save(model, f"CIFAR_model_{i}.pth")
    print("模型已保存！")

writer.close()

# tensorboard命令：tensorboard --logdir=CIFAR_logs --port=6007

输出：

Files already downloaded and verified
Files already downloaded and verified
训练数据集的长度：50000
测试数据集的长度：10000
--------第0轮训练开始--------
训练次数：100，Loss：2.289702892303467
训练次数：200，Loss：2.27992582321167
训练次数：300，Loss：2.2398457527160645
训练次数：400，Loss：2.1376163959503174
训练次数：500，Loss：2.0379090309143066
训练次数：600，Loss：1.986525058746338
训练次数：700，Loss：2.0220046043395996
整体测试集上的Loss：316.75445556640625
整体测试集上的正确率：0.2759000062942505
模型已保存！
--------第1轮训练开始--------
训练次数：800，Loss：1.8803398609161377
训练次数：900，Loss：1.8569337129592896
训练次数：1000，Loss：1.9385566711425781
训练次数：1100，Loss：1.9486432075500488
训练次数：1200，Loss：1.712395191192627
训练次数：1300，Loss：1.6654386520385742
训练次数：1400，Loss：1.749714732170105
训练次数：1500，Loss：1.8054516315460205
整体测试集上的Loss：296.45361328125
整体测试集上的正确率：0.32420000433921814
模型已保存！
--------第2轮训练开始--------
训练次数：1600，Loss：1.7527278661727905
.....
.....

TensorBoard输出：

补充

计算测试集上的正确率

argmax()：参数说明

代码实现：

import torch


output = torch.tensor(   [[0.1, 0.2],
                         [0.05, 0.4]]  )

print(output.argmax(1))  # tensor([1, 1])
print(output.argmax(0))  # tensor([0, 1])

preds = output.argmax(1)  # 1.取得在 1 方向，概率最大的索引，即得标签输出值
targets = torch.tensor([0, 1])
print(preds == targets)
print((preds == targets).sum())  # 2.对比标签输出值与目标值+求和：True=1，False=0

输出：

tensor([1, 1])
tensor([0, 1])
tensor([False,  True])
tensor(1)