pytorch入门21：在gpu上训练模型

"""
网络模型，数据(输入，target)，损失函数可以使用gpu跑，直接调用.cuda()即可

可以尝试使用Google Colab——用谷歌免费GPU跑你的深度学习代码
"""

import torch
import torchvision
from torch import nn
from torch.nn import Conv2d, MaxPool2d, Flatten, Linear

from torch.utils.data import DataLoader
# 1.准备数据集并加载
from torch.utils.tensorboard import SummaryWriter

train = torchvision.datasets.CIFAR10('./torchvision_dataset',train=True,download=False,transform=torchvision.transforms.ToTensor())
test = torchvision.datasets.CIFAR10('./torchvision_dataset',train=False,download=False,transform=torchvision.transforms.ToTensor())
train_loader = DataLoader(train,batch_size=64,drop_last=True)
test_loader = DataLoader(test,batch_size=64,drop_last=True)

#获取两个数据集的长度
train_len = len(train)
test_len = len(test)

# 定义训练的设备
#device = torch.device('cuda:0')
device = torch.device('cpu')

#定义网络结构
class Zkl(nn.Module):
    def __init__(self):
        super(Zkl, self).__init__()
        # self.conv1 = Conv2d(3,32,5,padding=2)
        # self.maxpool1 = MaxPool2d(2)
        # self.conv2 = Conv2d(32,32,5,padding=2)
        # self.maxpool2 = MaxPool2d(2)
        # self.conv3 = Conv2d(32,64,5,padding=2)
        # self.maxpool3 = MaxPool2d(2)
        # self.flatten = Flatten()
        # self.linear1 = Linear(1024,64)
        # self.linear2 = Linear(64,10)

        self.modle1 = nn.Sequential(
            Conv2d(3, 32, 5, padding=2),
            MaxPool2d(2),
            Conv2d(32, 32, 5, padding=2),
            MaxPool2d(2),
            Conv2d(32, 64, 5, padding=2),
            MaxPool2d(2),
            Flatten(),
            Linear(1024, 64),
            Linear(64,10)
        )

    def forward(self,x):
        # x = self.conv1(x)
        # x = self.maxpool1(x)
        # x = self.conv2(x)
        # x = self.maxpool2(x)
        # x = self.conv3(x)
        # x = self.maxpool3(x)
        # x = self.flatten(x)
        # x = self.linear1(x)
        # x = self.linear2(x)
        x = self.modle1(x)
        return x

#定义模型对象
zkl = Zkl()
#放在gpu上跑，调用cuda
# if torch.cuda.is_available():
#     zkl = zkl.cuda()
zkl = zkl.to(device)  # 将我们的网络转移到设备上面

#定义损失函数
loss = nn.CrossEntropyLoss()
#放在gpu上跑，调用cuda
# if torch.cuda.is_available():
#     loss = loss.cuda()
loss = loss.to(device)

#定义优化器
learn_rate = 1e-2
optim = torch.optim.SGD(zkl.parameters(),lr=learn_rate)

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

#定义训练轮数
epoch=10

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

for i in range(epoch):
    print(f"-------第 {i} 轮训练开始----------")
    zkl.train() # 设置这个是使模型进入一个特定状态，只对dropout层或者batchnorml层有作用，并且有这些层就必须调用
    for data in train_loader:
        imgs,targets = data
        # if torch.cuda.is_available():
        #     imgs = imgs.cuda()
        #     targets = targets.cuda()
        imgs = imgs.to(device)
        targets = targets.to(device)
        output = zkl(imgs)
        train_loss=loss(output,targets)
        #优化器调优模型
        optim.zero_grad()
        train_loss.backward()
        optim.step()

        total_train_step += 1
        if total_train_step%100 == 0:  # 减少打印次数，防止看不清楚结果
            print(f"训练次数 {total_train_step}：loss={train_loss.item()}")  # 加上.item()是为了使输出的是一个值，而不是tensor(1)这样的表达形式
            writer.add_scalar('model_train',train_loss,total_train_step)
    #每一轮训练完后开始测试
    zkl.eval()  # 设置这个是使模型进入一个特定状态，只对dropout层或者batchnorml层有作用，并且有这些层就必须调用
    total_test_loss = 0
    total_accuracy = 0
    with torch.no_grad():  # 不进行优化器在测试过程中，只需要使用训练好的参数
        for data in test_loader:
            imgs,targets = data
            # if torch.cuda.is_available():
            #     imgs = imgs.cuda()
            #     targets = targets.cuda()
            imgs = imgs.to(device)
            targets = targets.to(device)
            output = zkl(imgs)
            loss_test = loss(output,targets)
            total_test_loss += loss_test.item()
            accuracy = (output.argmax(1) == targets).sum()
            total_accuracy += accuracy

    print(f"第 {i} 轮训练后测试的loss: {total_test_loss}")
    print(f"第 {i} 轮训练后测试的accuracy: {total_accuracy/test_len}")
    writer.add_scalar('model_test', loss_test, total_test_step)
    writer.add_scalar('model_accuracy',total_accuracy/test_len,total_test_step)
    total_test_step += 1

    #保存每一轮训练的结果
    torch.save(zkl,f'modle_train_{i}.pth')
    print('模型已保存')

writer.close()