模型保存与读取

完整模型训练套路

import torch
import torchvision.datasets
from torch import nn
from torch.nn import Conv2d, MaxPool2d, Flatten, Linear
from torch.utils.data import DataLoader
from torch.utils.tensorboard import SummaryWriterfrom model import *train_data=torchvision.datasets.CIFAR10("data_nn",train=True,transform=torchvision.transforms.ToTensor(),download=True )
test_data=torchvision.datasets.CIFAR10("data_nn",train=False,transform=torchvision.transforms.ToTensor(),download=True )
train_data_size=len(train_data)
test_data_size=len(test_data)
print(test_data_size,train_data_size)train_dataloader=DataLoader(train_data,batch_size=64)
test_dataloader=DataLoader(test_data,batch_size=64)
#搭建神经网络XKK=xkkk()
#损失函数
loss_fn=nn.CrossEntropyLoss()
#优化器
learning_rate=0.01
optimizer=torch.optim.SGD(XKK.parameters(),lr=learning_rate)
#设置训练网络的一些参数
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=dataoutputs=XKK(imgs)loss=loss_fn(outputs,targets)#优化器调优，优化模型optimizer.zero_grad()loss.backward()optimizer.step()total_test_step=total_test_step+1if total_test_step%100==0:print("训练次数：{}，loss:{}".format(total_test_step,loss))writer.add_scalar("train_loss",loss.item(),total_test_step)#测试步骤开始total_test_loss=0with torch.no_grad():for data in test_dataloader:imgs,targets=dataoutputs=XKK(imgs)loss=loss_fn(outputs,targets)total_test_loss=total_test_loss+lossprint("整体测试集上的Loss:{}".format(total_test_loss))writer.add_scalar("test_loss",total_test_loss,total_test_step)total_test_step=total_test_step+1
writer.close()

使用GPU训练

import timeimport torch
import torchvision.datasets
from torch import nn
from torch.nn import Conv2d, MaxPool2d, Flatten, Linear
from torch.utils.data import DataLoader
from torch.utils.tensorboard import SummaryWriter# from model import *train_data=torchvision.datasets.CIFAR10("data_nn",train=True,transform=torchvision.transforms.ToTensor(),download=True )
test_data=torchvision.datasets.CIFAR10("data_nn",train=False,transform=torchvision.transforms.ToTensor(),download=True )
train_data_size=len(train_data)
test_data_size=len(test_data)
print(test_data_size,train_data_size)train_dataloader=DataLoader(train_data,batch_size=64)
test_dataloader=DataLoader(test_data,batch_size=64)
#搭建神经网络class xkkk(torch.nn.Module):def __init__(self):super(xkkk,self).__init__()self.model1=torch.nn.Sequential(Conv2d(3, 32, 5,1, padding=2),MaxPool2d(kernel_size=2),Conv2d(32, 32, 5, 1,padding=2),MaxPool2d(kernel_size=2),Conv2d(32, 64, 5,stride=1, padding=2),MaxPool2d(kernel_size=2),Flatten(),Linear(in_features=64*4*4, out_features=64),  # 1024=64*4*4,Linear(64, 10))def forward(self, x):output=self.model1(x)return output
XKK=xkkk()
XKK=XKK.cuda()
#损失函数
loss_fn=nn.CrossEntropyLoss()
loss_fn=loss_fn.cuda()
#优化器
learning_rate=0.01
optimizer=torch.optim.SGD(XKK.parameters(),lr=learning_rate)
#设置训练网络的一些参数
total_test_step=0
#训练的轮数
epoch=10
#添加tensorboard
writer=SummaryWriter("logs_train")
start_time=time.time()
for i in range(epoch):print("----第{}轮训练开始----".format(i+1))for data in train_dataloader:imgs,targets=dataimgs=imgs.cuda()targets=targets.cuda()outputs=XKK(imgs)loss=loss_fn(outputs,targets)#优化器调优，优化模型optimizer.zero_grad()loss.backward()optimizer.step()total_test_step=total_test_step+1if total_test_step%100==0:end_time=time.time()print(end_time-start_time)print("训练次数：{}，loss:{}".format(total_test_step,loss))writer.add_scalar("train_loss",loss.item(),total_test_step)#测试步骤开始total_test_loss=0with torch.no_grad():for data in test_dataloader:imgs,targets=dataimgs = imgs.cuda()targets = targets.cuda()outputs=XKK(imgs)loss=loss_fn(outputs,targets)total_test_loss=total_test_loss+loss.item()print("整体测试集上的Loss:{}".format(total_test_loss))writer.add_scalar("test_loss",total_test_loss,total_test_step)total_test_step=total_test_step+1torch.save(XKK,"XKK_{}.pth".format(i))print("模型已保存")
writer.close()

CPU训练时间如下

使用GPU训练时间如下

对比可知GPU训练速度提升5倍左右

GPU代码区CPU区别如下

只需要改动网络模型、数据、损失函数，调用它们对应的.cuda()。

方式2

#定义训练的设备
#device=torch.device("cpu")#使用cpu
device=torch.device("cuda")#使用GPU
....
XKK=xkkk()
XKK=XKK.to(device)
......imgs=imgs.to(device)targets=targets.to(device).....

测试

import torch
import torchvision.transforms
from PIL import Image
from torch.nn import Conv2d, MaxPool2d, Flatten, Linearimage_path="imgs/dog.png"image = (Image.open(image_path).convert("RGB"))
#image=Image.open(image_path)
transform=torchvision.transforms.Compose([torchvision.transforms.Resize((32,32)),torchvision.transforms.ToTensor()])
image=transform(image)class xkkk(torch.nn.Module):def __init__(self):super(xkkk,self).__init__()self.model1=torch.nn.Sequential(Conv2d(3, 32, 5,1, padding=2),MaxPool2d(kernel_size=2),Conv2d(32, 32, 5, 1,padding=2),MaxPool2d(kernel_size=2),Conv2d(32, 64, 5,stride=1, padding=2),MaxPool2d(kernel_size=2),Flatten(),Linear(in_features=64*4*4, out_features=64),  # 1024=64*4*4,Linear(64, 10))def forward(self, x):output=self.model1(x)return output
model=torch.load("XKK_9.pth",map_location=torch.device("cpu"))
print(model)
print(image.size())image=torch.reshape(image,(1,3,32,32))
model.eval()
with torch.no_grad():output=model(image)
print(output)
print(output.argmax(1))

输入一张小狗图片

image_path="imgs/dog.png"结果
tensor([5])

输入一张飞机图片

image_path="imgs/airplane.png"
结果
tensor([0])

可以看出预测准确