利用GPU进行训练

文章目录

一、GPU训练模型

GPU只能够训练三种变量，分别是：

网络模型
数据（输入，标注targets）
损失函数

使用方式是.cuda()

除了以上的三种，在别的地方是引用不上cuda的。

使用方式：

1.网络模型

bash 复制代码

#创建网络模型
class Sen(nn.Module):
    def __init__(self):
        super(Sen, self).__init__()
        self.model = nn.Sequential(
            nn.Conv2d(3, 32, 5, 1, 2),
            nn.MaxPool2d(2),
            nn.Conv2d(32, 32, 5, 1 ,2),
            nn.MaxPool2d(2),
            nn.Conv2d(32, 64, 5, 1, 2),
            nn.MaxPool2d(2),
            nn.Flatten(),
            nn.Linear(64*4*4, 64),
            nn.Linear(64, 10)
        )
    def forward(self,x):
        x = self.model(x)
        return x
sen = Sen()
sen.cuda()

2、损失函数

bash 复制代码

#损失函数
loss_fn = nn.CrossEntropyLoss()
loss_fn = loss_fn.cuda()

3、数据

bash 复制代码

    for data in train_dataloader:
        imgs, targets = data
        imgs = imgs.cuda()
        targets = targets.cuda()
        outputs = sen(imgs)
        loss = loss_fn(outputs, targets)

如果电脑上没有GPU的话，以上这么写会发生报错，更好的写法应该是加上if判断是否有GPU

例下：

bash 复制代码

if torch.cuda.is_available():
    sen.cuda()

bash 复制代码

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

bash 复制代码

if torch.cuda.is_available():
  imgs = imgs.cuda()
  targets = targets.cuda()

这样可以保证有gpu先试用gpu，没有gpu也能够正常运行。

二、对比使用gpu和cpu进行训练所花费的时间

引入记录时间函数：

bash 复制代码

import time

bash 复制代码

start_time = time.time()

bash 复制代码

end_time = time.time()

bash 复制代码

print(f"使用gpu训练使用的时间{end_time - start_time}")

在模型训练开始与结束时候分别记录start_time和end_time:

bash 复制代码

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

    #训练步骤开始
    sen.train()
    for data in train_dataloader:
        imgs, targets = data
        imgs = imgs.cuda()
        targets = targets.cuda()
        outputs = sen(imgs)
        loss = loss_fn(outputs, targets)

        #优化器模型
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        total_train_step = total_train_step + 1
        if total_train_step % 100 == 0:
            end_time = time.time()
            print(f"使用gpu训练使用的时间{end_time - start_time}")
            print(f"训练次数：{total_train_step}，Loss：{loss.item()}")
            writer.add_scalar("train_loss", loss.item(), total_train_step)

1.使用gpu训练结果：

可以看到平均100次只使用1s。

2.使用cpu进行训练结果：

可以看到没训练100次花费时间是7s

因此使用gpu进行训练模型能大大提升算力。

三、GPU训练模型的第二种表达方式

先定义训练的设备：

bash 复制代码

device = torch.device("cpu")

再在变量当中调用

bash 复制代码

sen = Sen()
sen = sen.to(device)

bash 复制代码

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

bash 复制代码

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

调用gpu只需要改变训练模型即可，变量当中依旧用.to(device)：

bash 复制代码

device = torch.device("cuda")

'：0'的写法等同于cuda

bash 复制代码

device = torch.device("cuda:0")

如果电脑上有两个不同的gpu，可以用以下调用第二个gpu进行训练

bash 复制代码

device = torch.device("cuda:1")

为了适用于不同环境，可以用以下代码来写：

bash 复制代码

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