模型保存与读取
完整模型训练套路
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 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)
#搭建神经网络
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=data
outputs=XKK(imgs)
loss=loss_fn(outputs,targets)
#优化器调优，优化模型
optimizer.zero_grad()
loss.backward()
optimizer.step()
total_test_step=total_test_step+1
if 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=0
with torch.no_grad():
for data in test_dataloader:
imgs,targets=data
outputs=XKK(imgs)
loss=loss_fn(outputs,targets)
total_test_loss=total_test_loss+loss
print("整体测试集上的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 time
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 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=data
imgs=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+1
if 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=0
with torch.no_grad():
for data in test_dataloader:
imgs,targets=data
imgs = 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+1
torch.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, Linear
image_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])
可以看出预测准确
