pytorch【3】完整的模型训练过程

1.模型训练套路

import torch
import torchvision
from torch.utils.data import DataLoader
from torch import nn
from torch.nn import Sequential, Conv2d, MaxPool2d, CrossEntropyLoss, Flatten, Linear
import time

from torch.utils.tensorboard import SummaryWriter


class My(nn.Module):
    def __init__(self):
        super(My, self).__init__()
        self.model = Sequential(
            Conv2d(3, 32, 5, padding=2),
            MaxPool2d(2),
            Conv2d(32, 32, 5, padding=2),
            MaxPool2d(2),
            Conv2d(32, 64, 5, padding=2),
            MaxPool2d(2),
            Flatten(),
            Linear(1024, 64),
            Linear(64, 10)
        )

    def forward(self, x):
        x = self.model(x)
        return x



# 格式处理
transform = torchvision.transforms.Compose([
    torchvision.transforms.ToTensor()
])
# 加载数据集
train_data = torchvision.datasets.CIFAR10(root='./dataset', train=True, transform=transform, download=True)
test_data = torchvision.datasets.CIFAR10(root='./dataset', train=False, transform=transform, download=True)
# 批量处理
train_dataloader = DataLoader(dataset=train_data, batch_size=64)
test_dataloader = DataLoader(dataset=test_data, batch_size=64)

train_len = len(train_data)
test_len = len(test_data)

# 实例化模型
mymodel = My()
mymodel = mymodel.cuda()

# 损失函数
loss = CrossEntropyLoss()
loss = loss.cuda()



# 优化器

leanr_rate = 1e-2
optim = torch.optim.SGD(mymodel.parameters(), lr=leanr_rate)


train_step = 0
test_step = 0
start_time = time.time()
writer = SummaryWriter("./logs_train")
for epoch in range(20):
    print("---------------第{}轮训练-------------".format(epoch))
    # 训练集
    mymodel.train()
    for data in train_dataloader:
        imgs, targets = data
        imgs = imgs.cuda()
        targets = targets.cuda()
        output = mymodel(imgs)
        loss_result = loss(output, targets)
        optim.zero_grad()
        loss_result.backward()
        optim.step()
        if train_step % 100 == 0:
            end_time = time.time()
            print(end_time - start_time)
            print("训练次数:{}, Loss:{}".format(train_step, loss_result.item()))
            writer.add_scalar("train_loss", loss_result.item(), train_step)
        train_step += 1

    # 测试集
    total_loss_result = 0
    total_accuracy = 0
    with torch.no_grad():
        for data in test_dataloader:
            imgs, targets = data
            imgs = imgs.cuda()
            targets = targets.cuda()
            output = mymodel(imgs)
            accuracy = (output.argmax(1) == targets).sum()
            loss_result = loss(output, targets)
            total_accuracy += accuracy
            total_loss_result += loss_result
    print("整体测试集上的Loss:{}".format(total_loss_result))
    print("整体测试集上的Accuracy:{}".format(total_accuracy/test_len))
    writer.add_scalar("test_loss", total_loss_result.item(), test_step)
    writer.add_scalar("accuracy", total_accuracy/test_len, test_step)
    test_step += 1

    torch.save(mymodel, "./models/mymodel_{i}.pth".format(i=epoch))
    print("模型已保存")

writer.close()

2.使用GPU

谷歌免费：https://colab.research.google.com/
网络模型、数据（输入和标注）、损失函数 ，调用.cuda()
法一：如果存在就使用GPU，否则CPU

if torch.cuda.is_available():
    mymodel = mymodel.cuda()
...
...
if torch.cuda.is_available():
    loss = loss.cuda()

法二：指定设备

# 指定设备
# device = torch.device("cuda")
# device = torch.device("cuda:0")
# device = torch.device("cpu")
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# 实例化模型
mymodel = My()
mymodel = mymodel.to(device)

# 损失函数
loss = CrossEntropyLoss()
loss = loss.to(device)

3.验证模型（就是应用）

提供一个图片、输出类型
特别注意：

# 1.需要引入网络模型 class My(nn.Module):XXX  -->这是因为保存时用的方法1
# 2.加载模型时weights_only=False必须的，
# 3.保存的模型是使用GPU训练出来的，需要加载时指明使用环境为cpu:map_location=torch.device('cpu')
class My(nn.Module):
    def __init__(self):
        super(My, self).__init__()
	XXXXXXXXXXXX
mymodel = torch.load('./models/mymodel_19.pth', weights_only=False, map_location=torch.device('cpu'))


# Conv2d的输入要求是四维的数据
image = torch.reshape(image, (1, 3, 32, 32))

# 小注意
mymodel.eval()
with torch.no_grad():
    output = mymodel(image)

import torch
from PIL import Image
import torchvision
from torch import nn
from torch.nn import Linear, Conv2d, MaxPool2d, Flatten, Sequential
image_path = "./image/gou.jpg"
image = Image.open(image_path)
transform = torchvision.transforms.Compose([
    torchvision.transforms.Resize((32, 32)),
    torchvision.transforms.ToTensor()
])
image = transform(image)
print(image.shape)
class My(nn.Module):
    def __init__(self):
        super(My, self).__init__()
        self.model = Sequential(
            Conv2d(3, 32, 5, padding=2),
            MaxPool2d(2),
            Conv2d(32, 32, 5, padding=2),
            MaxPool2d(2),
            Conv2d(32, 64, 5, padding=2),
            MaxPool2d(2),
            Flatten(),
            Linear(1024, 64),
            Linear(64, 10)
        )

    def forward(self, x):
        x = self.model(x)
        return x

mymodel = torch.load('./models/mymodel_19.pth', weights_only=False, map_location=torch.device('cpu'))

image = torch.reshape(image, (1, 3, 32, 32))

mymodel.eval()
with torch.no_grad():
    output = mymodel(image)
print(output.argmax(1).item())
print(output)