【Pytorch 入门】DAY5 完整的模型训练套路与模型验证

完整的模型训练套路

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

from model import *

#准备数据集
train_data=torchvision.datasets.CIFAR10('dataset',train=True,transform=torchvision.transforms.ToTensor(),download=True)
test_data=torchvision.datasets.CIFAR10('dataset',train=False,transform=torchvision.transforms.ToTensor(),download=True)
#查看数据集长度
train_data_size=len(train_data)
test_data_size=len(test_data)
# print("训练数据集长度为：{}".format(train_data_size))#训练数据集长度为：50000
# print("测试数据集长度为：{}".format(test_data_size))#测试数据集长度为：10000

#利用dataloader来加载数据
train_dataloader=DataLoader(train_data,batch_size=64)
test_dataloader=DataLoader(test_data,batch_size=64)

#创建网络模型
felix=Felix()

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

#优化器
learning_rate=1e-2#相当于 learning_rate=0.01
optimizer=torch.optim.SGD(felix.parameters(),lr=learning_rate)

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

#添加tensorboard
writer=SummaryWriter('logs_train')

for i in range(epoch):
    print("---------第{}轮训练开始--------".format(i+1))
    #训练步骤开始
    for data in train_dataloader:
        imgs,targets=data
        outputs=felix(imgs)
        loss=loss_fn(outputs,targets)

        #优化器优化模型
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        total_train_step=total_train_step+1
        if total_train_step%100==0:
            print('训练次数：{},loss:{}'.format(total_train_step, loss.item()))
            writer.add_scalar('train_loss',loss.item(),total_train_step)


    #测试步骤开始
    total_test_loss=0
    total_accuracy=0
    with torch.no_grad():
        for data in test_dataloader:
            imgs,targets=data
            outputs=felix(imgs)
            loss=loss_fn(outputs,targets)
            total_test_loss=total_test_step+loss
            accuracy=(outputs.argmax(1)==targets).sum()
            total_accuracy=total_accuracy+accuracy

    print('整体测试集上的loss:{}'.format(total_test_loss))
    print('整体测试集上的正确率：{}'.format(total_accuracy/test_data_size))
    writer.add_scalar('test_loss',total_test_loss,total_test_step)
    writer.add_scalar('test_accuracy',total_accuracy/test_data_size,total_train_step)
    total_test_step=total_test_step+1

    torch.save(felix,'felix_{}.pth'.format(i))
    print("模型已保存")

writer.close()

#搭建神经网络,放入单独的文件中
import torch
from torch import nn


class Felix(nn.Module):
    def __init__(self):
        super(Felix,self).__init__()
        self.model=nn.Sequential(
            nn.Conv2d(3,32,5,1,2),
            nn.MaxPool2d(2),
            nn.Conv2d(32,32,5,1,2),
            nn.MaxPool2d(2),
            nn.Conv2d(32,64,5,1,2),
            nn.MaxPool2d(2),
            nn.Flatten(),
            nn.Linear(64*4*4,64),
            nn.Linear(64,10)
        )
    def forward(self,x):
        x=self.model(x)
        return x

if __name__=='__main__':
    felix=Felix()
    input=torch.ones((64,3,32,32))
    output=felix(input)
    print(output.shape)#torch.Size([64, 10])

运行结果

模型验证

利用已训练好的模型，然后给它提供输入

用于预测的图像1

使用第一次的训练模型

import torch
import torchvision
from PIL import Image
from torch import nn

image_path='images/img.png'
image=Image.open(image_path)
print(image)
image=image.convert('RGB')

transform=torchvision.transforms.Compose([torchvision.transforms.Resize((32,32)),
                                          torchvision.transforms.ToTensor()])
image=transform(image)
print(image.shape)

class Felix(nn.Module):
    def __init__(self):
        super(Felix,self).__init__()
        self.model=nn.Sequential(
            nn.Conv2d(3,32,5,1,2),
            nn.MaxPool2d(2),
            nn.Conv2d(32,32,5,1,2),
            nn.MaxPool2d(2),
            nn.Conv2d(32,64,5,1,2),
            nn.MaxPool2d(2),
            nn.Flatten(),
            nn.Linear(64*4*4,64),
            nn.Linear(64,10)
        )
    def forward(self,x):
        x=self.model(x)
        return x

if __name__=='__main__':
    felix=Felix()
    input=torch.ones((64,3,32,32))
    output=felix(input)
    print(output.shape)#torch.Size([64, 10])

#第一种方式保存，第一种方式加载
model=torch.load('felix_0.pth')
print(model)
image=torch.reshape(image,(1,3,32,32))
model.eval()
with torch.no_grad():
    output=model(image)

print(output)

print(output.argmax(1))

结果不准，因为模型没怎么训练

使用第九次训练的模型

model=torch.load('felix_9.pth')

预测正确

用于预测的图像2

预测正确，是船

完结散花！！！！