Pytorch 学习笔记06 完整的模型训练套路，GPU使用，完整模型的验证套路，如何计时

Pytorch 学习笔记06

完整的模型训练套路（以CIFAR10数据集为例）

1. 导入数据集，数据加载器
   导入torchvision的官方数据集CIFAR10，设置训练集和测试集，训练数据加载器和测试数据加载器

2. 搭建神经网络模型
   利用CIFAR10模型搭建
   

3. 设置Loss函数
   CIFAR10数据集是分类问题，所以使用nn.CrossEntropyLoss()

4. 设置优化函数
   使用torch.optim.SGD（）

5. 设置训练网络的一些参数（不是非必要）
   训练的次数，测试的次数，训练的轮数

6. 开始每一轮训练+测试

   提问： 每次训练后怎么知道模型是不是最优？
   回答： 每次训练后都在测试集上进行测试

   1. 将测试集在模型上跑，并将梯度设置为零，避免测试集优化模型
   2. 计算每一轮中，测试集在模式行跑的总Loss，记录正确数据的个数（用来计算每轮测试集的正确率）
   3. 优化：通过tensorboard展示训练集每一批的Loss；测试集每一轮的Loss、正确率。
   4. 保存每一轮的模型

代码：

import torch
import torchvision
from torch.utils.tensorboard import SummaryWriter

from model_CIFAR10 import *

# 1导入数据集
from torch.utils.data import DataLoader

train_dataset = torchvision.datasets.CIFAR10("./dataset/CIFAR10", train=True,
                                             transform=torchvision.transforms.ToTensor(),
                                             download=True)
test_dataset = torchvision.datasets.CIFAR10("./dataset/CIFAR10", train=False,
                                             transform=torchvision.transforms.ToTensor(),
                                             download=True)

# 1 导入数据加载器
train_dataloader = DataLoader(train_dataset,batch_size=64)
test_dataloader = DataLoader(test_dataset,batch_size=64)

# 2 创建神经网络
mynn = nn_CIFAR()

# 3 loss函数
loss_func = nn.CrossEntropyLoss()

# 4 优化函数
learn_rate = 1e-2 # 学习速率
optim = torch.optim.SGD(mynn.parameters(), lr = learn_rate)

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

# 开始训练

writer = SummaryWriter("logs_CIFAR10")

for i in range(epoch):
    print(f"------第{i+1}轮训练开始------")
    #开始训练，
    mynn.train() #表示模型在训练模式，
    # train()只对一些特定层有用，可以看官方文档，以后的模型中有特殊的层就必须加上train()，这里没用
    for imgs, tars in train_dataloader:
        outputs = mynn(imgs) #放入模型
        loss = loss_func(outputs,tars) # 计算loss
        optim.zero_grad() # 梯度清零
        loss.backward() # 计算梯度
        optim.step() # 更新参数
        total_train_step += 1
        if total_train_step%100 ==0 : # 每训练一百次进行展示
            print(f"训练次数为{total_train_step}，Loss：{loss}")
            writer.add_scalar("train_loss", loss.item(), total_train_step) # 每百次的训练loss

    """
    提问： 每次训练后怎么知道模型是不是最优？
    回答： 每次训练后都在测试集上进行测试
    """
    # 进行测试
    mynn.eval() # 表示模型在测试模式，与train() 同理
    total_test_loss = 0.0 #记录整体测试集的loss
    total_accuracy = 0 # 记录每一轮正确的测试数据
    with torch.no_grad():   # 将梯度设为0，保证模型在测试集上跑不会对模型进行调优。只需要对模型进行测试，不需要优化
        for imgs, tars in test_dataloader:
            outputs = mynn(imgs)
            loss = loss_func(outputs, tars)
            total_test_loss += loss
            accuracy = (outputs.argmax(1) == tars).sum() #记录每一批数据中正确的个数
            """
            outputs为tensor类型数据
            argmax(x)函数为求一行或一列的最大值，参数x为1则是求一行，为0则是求一列
            将outputs.argmax(1) == tars相比，相等的为True，不等的为False
            然后对结果求和，True为1，False为0，即为正确的个数
            """
            total_accuracy += accuracy # 求和
        print(f"第{i+1}轮训练，整体测试集上的Loss：{total_test_loss}")
        print(f"第{i+1}轮训练，整体测试集上的正确率：{total_accuracy/len(test_dataset)}")
        writer.add_scalar("test_loss", total_test_loss, i+1) # 每一轮测试集总loss
        writer.add_scalar("test_accuracy", total_accuracy/len(test_dataset), i+1) # 每一轮测试集正确率

    # 保存每一轮的模型
    torch.save(mynn, f"./CIFAR_train_model/CIFAR_{i+1}.pth")
    print(f"第{i+1}轮模型已保存")

writer.close()

模型：

import torchvision
from torch import nn
from torch.utils.data import DataLoader



class nn_CIFAR(nn.Module):
    def __init__(self):
        super(nn_CIFAR, self).__init__()
        self.model = nn.Sequential(
            nn.Conv2d(3, 32, kernel_size=5, padding= 2), # stride使用默认值1，padding是根据官方文档里的公式计算得出
            nn.MaxPool2d(2),
            nn.Conv2d(32, 32, kernel_size=5, padding=2),
            nn.MaxPool2d(2),
            nn.Conv2d(32, 64, kernel_size=5, padding=2),
            nn.MaxPool2d(2),
            nn.Flatten(),
            nn.Linear(1024, 64),  # 1024 为 64 * 4 *4 得出
            nn.Linear(64, 10)
        )
    def forward(self,x):
        return self.model(x)


if __name__ == '__main__':
    test = nn_CIFAR()

利用GPU训练的两种方式

对以下三个对象进行GPU的调用

1. 在三个对象后用 .cuda()进行GPU调用

# 神经网络模型对象
mynn = nn_CIFAR() 
if torch.cuda.is_available():
    mynn = mynn.cuda()
# Loss函数
loss_func = nn.CrossEntropyLoss()
if torch.cuda.is_available():
    loss_func = loss_func.cuda()
# 数据
if torch.cuda.is_available():
     imgs = imgs.cuda()
     tars = tars.cuda()

2. 通过torch.device()，指定设备。

torch.device("cpu")  # 指定为CPU
torch.device("cuda")  # 指定为GPU
# 若有多个GPU
torch.device("cuda:0")  # 指定为第一个GPU
torch.device("cuda:0")  # 指定为第二个GPU

#定义训练设备
device = torch.device("cuda")
# 神经网络模型对象
mynn = nn_CIFAR()
mynn = mynn.to(device)
# Loss函数
loss_func = nn.CrossEntropyLoss()
loss_func = loss_func.to(device)
# 数据
imgs = imgs.to(device)
tars = tars.to(device)

完整模型的验证套路

"""
测试CIFAR10模型

"""
import torch
import torchvision
from PIL import Image
from model_CIFAR10 import * # import模型代码
# 随机找的图片
img_path = "imgs/airplane.png" 
# 导入图片
img = Image.open(img_path)
# 设置transforms,将图片转为tensor
transform = torchvision.transforms.Compose([torchvision.transforms.Resize((32,32)),
                                      torchvision.transforms.ToTensor()])

img = transform(img)
print(img.shape)
# 加载训练好的模型
model = torch.load("CIFAR_30.pth", map_location=torch.device('cpu'))

img = torch.reshape(img,(1,3,32,32))
# 模型测试模式
model.eval()
with torch.no_grad():
    output = model(img)
# 找出得分最高
output = output.argmax(1)
ans = ['airplane', 'automobile', 'bird', 'cat', 'deer', 'dog', 'frog', 'horse', 'ship',
       'truck']
print(ans[output[0].item()])

如何计时

import time
start_time = time.time()
"""
一系列代码
"""
end_time = time.time()
print(end_time - start_time)