18.train_GPU

方式一

gpu通常使用在:

模型 demo = demo.cuda()

损失函数l;oss_fn = loss_fn.cuda()

数据：[训练],[验证]

        imgs = imgs.cuda()
        targets = targets.cuda()

通常是使用这种方法

if torch.cuda.is_available():
    demo = demo.cuda()

经测试：

在使用cpu的情况下，total_train_step = 100 所使用的时间为10.69461178779602

在使用gpu的情况下，total_train_step = 100 所使用的时间为3.9713311195373535

方式二

在开头加上一下两种形式的代码：

对于单显卡,以下两种写法没有区别。

device =  torch.device("cuda")
device =  torch.device("cuda:0") # 多张显卡时，可以用来指定显卡 0代表第一张显卡

这种方式更加常见

device =  torch.device("cuda" if torch.cuda.is_available() else "cpu")

tips:

无论是哪种使用gpu的方式，在model 和 loss 处都可以不用重新赋值，但在数据（例如图片和标签处）必须重新赋值。

例如

# 模型
demo.to(device)
demo.cuda()
# loss
loss_fn.to(device)
loss_fn.cuda()
# imgs and targets
imgs = imgs.to(device)
targets = targets.to(device)
if torch.cuda.is_available():
    imgs = imgs.cuda()
    targets = targets.cuda()

code：

import torch
import torchvision
from torch.utils.data import DataLoader
from torch.utils.tensorboard import SummaryWriter
#import time
from models import *
# 定义训练时使用的设备
#device =  torch.device("cpu")
#device =  torch.device("cuda")
device =  torch.device("cuda" if torch.cuda.is_available() else "cpu")
# 准备数据集
train_data = torchvision.datasets.CIFAR10(root="./dataset",train=True,transform=torchvision.transforms.ToTensor(),
                                          download=True)
test_data = torchvision.datasets.CIFAR10(root="./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))
print("测试数据长度为：{}".format(test_data_size))
# 加载数据集
train_dataloader = DataLoader(train_data,batch_size=64)
test_dataloader = DataLoader(test_data,batch_size=64)
# 创建神经网络
demo = DEMO()
demo.to(device)
# demo = demo.to(device) 也可以的。第一种方式也是一样，但 图片 和 标签之类的数据必须重新赋值
# if torch.cuda.is_available():
#     demo = demo.cuda()
# 定义损失函数
loss_fn = nn.CrossEntropyLoss()
loss_fn.to(device)
#loss_fn = loss_fn.to(device) 也可以。第一种方式也是一样
# if torch.cuda.is_available():
#     loss_fn = loss_fn.cuda()
# 定义优化器
#learing_rate = 1e-2 = 1 × (10) ^ (-2)
learning_rate = 0.01
optimizer = torch.optim.SGD(demo.parameters(),lr=learning_rate)
# tensorboard
writer = SummaryWriter("./train_logs")
# 训练
# 设置一些参数
total_train_step = 0
total_test_step = 0
epoch = 10
#start_time = time.time()
for i in range(epoch):
    print("--------------epoch:{}----------------".format(i+1))
    # 训练
    demo.train() # 在使用一些特殊层时要调用，不适用时也可以调用，实践中很常用 同 demo.eval()
    for data in train_dataloader:
        imgs, targets = data
        # 图片 标签之类的数据 在使用cuda是必须要重新赋值
        imgs = imgs.to(device)
        targets = targets.to(device)
        # if torch.cuda.is_available():
        #     imgs = imgs.cuda()
        #     targets = targets.cuda()
        output = demo(imgs)
        loss = loss_fn(output,targets)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        total_train_step += 1
        if total_train_step % 100 == 0:
            #end_time = time.time()
            #print("total_train_step = 100 所使用的时间为{}".format(end_time - start_time))
            print("训练次数：{}，train_Loss = {}".format(total_train_step,loss.item()))
            writer.add_scalar("train_loss",loss.item(),total_train_step)
    # 验证
    demo.eval()
    total_test_loss = 0
    # 梯度为零，不优化参数
    total_acc = 0
    with torch.no_grad():
        for data in test_dataloader:
            imgs,targets = data
            imgs = imgs.to(device)
            targets = targets.to(device)
            # if torch.cuda.is_available():
            #     imgs = imgs.cuda()
            #     targets = targets.cuda()
            output = demo(imgs)
            loss = loss_fn(output, targets)
            total_test_loss = total_test_loss + loss
            pre = output.argmax(1)
            acc = ((pre == targets).sum())
            total_acc = total_acc + acc
    accuracy = total_acc / test_data_size
    print("total_test_loss = {}".format(total_test_loss))
    print("accuracy = {}".format(accuracy))
    writer.add_scalar("test_loss",total_test_loss,i+1)
    writer.add_scalar("accuracy",accuracy,i+1)
    torch.save(demo,"pretrained_demo_{}.pth".format(i+1))
    print("pretrained_demo已保存")
writer.close()