【Pytorch】使用GPU进行训练

【Pytorch】使用GPU进行训练

定义device

• device = torch.device(“cuda”)
  

多个GPU可以指定cuda的序号

import torchvision
from torch.utils.data import DataLoader
from torch import nn
import torch

# 定义训练的设备 使用gpu进行训练
device = torch.device("cuda")

# 准备数据集
train_data = torchvision.datasets.CIFAR10(root = "../data",train=True,transform=torchvision.transforms.ToTensor(),download=True)

test_data = torchvision.datasets.CIFAR10(root = "../data",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))

#  dataloader加载数据集
train_dataloader = DataLoader(train_data,batch_size = 64)
test_dataloader = DataLoader(test_data,batch_size = 64)

# 搭建神经网络

class Tudui(nn.Module):
    def __init__(self):
        super(Tudui,self).__init__()
        self.model = nn.Sequential(
            #  inpuT 3 output 32  卷积核5 stride 1 padding 2
            #  第一次卷积操作
            nn.Conv2d(3,32,5,1,2),

            #  最大池化操作 2 是kernel size
            nn.MaxPool2d(2),

            # 卷积操作 
            nn.Conv2d(32,32,5,1,2),

            nn.MaxPool2d(2),

            nn.Conv2d(32,64,5,1,2),

            nn.MaxPool2d(2),

            # 展平 将64个channel的4 x 4的图片全部展平 
            # 展平之后的结果是 6 4 * 4 * 4  一个一维向量
            nn.Flatten(),
            #  线性层
            nn.Linear(64 * 4 * 4,64),
            nn.Linear(64,10)
        )

    #  前向传播
    def forward(self,x):
        x = self.model(x)
        return x

import torch

#  创建网络模型
tudui = Tudui()
tudui = tudui.to(device)

# 定义损失函数
loss_fn = nn.CrossEntropyLoss()
loss_fn = loss_fn.to(device)

learning_rate = 0.01
# 定义优化器  随机梯度下降
optimizer = torch.optim.SGD(tudui.parameters(),lr = learning_rate)


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

# 训练的轮数
epoch = 10
from torch.utils.tensorboard import SummaryWriter

#  添加tensorboard

writer = SummaryWriter("../logs_train")

for i in range(epoch):
    print("------第{}轮训练开始".format(i + 1))

    # 训练步骤开始
    tudui.train()
    for data in train_dataloader:
        imgs,targets = data

        imgs = imgs.to(device)
        targets = targets.to(device)

        # 64 x 10
        outputs = tudui.forward(imgs)

        #  计算损失
        loss = loss_fn(outputs,targets)

        #  将梯度清0
        optimizer.zero_grad()

        # 反向传播  计算梯度
        loss.backward()

        # 优化
        optimizer.step()

        #  统计训练次数
        total_train_step = total_train_step + 1

        if total_train_step % 100 == 0:
            # item 将tensor转换为一个数据类型
            print("训练次数:{},Loss:{}".format(total_train_step,loss.item()))
            #  绘制训练损失
            writer.add_scalar("train_loss",loss.item(),total_train_step)

    #  测试步骤开始
    tudui.eval()
    total_test_loss = 0
    total_accuracy = 0
    with torch.no_grad():

        #  取出测试数据集的数据
        for data in test_dataloader:

            #  取出数据
            imgs,targets = data
            imgs = imgs.to(device)
            targets = targets.to(device)

            outputs = tudui(imgs)

            loss = loss_fn(outputs,targets) # 计算损失

            #  统计测试集上面的总损失
            total_test_loss = total_test_loss + loss.item()
            accuracy = (outputs.argmax(1) == targets).sum()
            total_accuracy = total_accuracy + accuracy


    print("整体测试集上面的Loss:{}".format(total_test_loss))
    print("整体测试及上面的正确率:{}".format(total_accuracy / test_data_size))
    writer.add_scalar("test_loss",loss.item(),total_test_step)
    writer.add_scalar("test_accuracy",total_accuracy / test_data_size,total_test_step)
    total_test_step = total_test_step + 1

使用cpu运行代码需要十一分钟，使用GPU运行代码需要2min多