【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
预测正确，是船
