动手学深度学习（pytorch土堆）-06损失函数与反向传播、模型训练、GPU训练

模型保存与读取

完整模型训练套路

c 复制代码

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训练

c 复制代码

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

c 复制代码

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

.....

测试

c 复制代码

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))

输入一张小狗图片

c 复制代码

image_path="imgs/dog.png"

结果
tensor([5])

输入一张飞机图片

c 复制代码

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

可以看出预测准确