Pytorch学习笔记(5) Pytorch GPU加速训练

主要介绍在Pytorch中如何使用GPU进行加速训练，同时在多GPU的条件下如何进行单机并行计算。

一、Pytorch使用GPU加速训练

常用命令

1、 获取设备里的GPU

#syntax 1
device = torch.device("cuda:0,1") 
#syntax 2  会根据设备情况自动选择cpu或者gpu
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

2、将模型和数据放到GPU中运行

# syntax 1
model.to(device)

# syntax 2 可以直接使用
model.gpu()
model.cuda()

# 将数据放入gpu
data = data.to(device)
data = data.gpu()
data = data.cuda()

通过以上例子知道如何将模型以及数据放入CPU中加速训练。其他的步骤基本保持不变。
当然，如果想使用cpu训练

device = torch.cuda("cpu")
model.to(device)
data.to(device)

# 或者
model.cpu()
data = data.cpu()

3、
在多 GPU 中执行前馈，后馈操作是非常自然的。尽管如此，PyTorch 默认只会使用一个 GPU。通过使用 DataParallel 让你的模型并行运行，你可以很容易的在多 GPU 上运行你的操作。

首先指定GPU运行的几种方式

# 终端指定
CUDA_VISIBLE_DEVICES=1 python script.py

# python 代码中指定
import os
os.environ["CUDA_VISIBLE_DEVICES"] = "0,1,2"
# 不使用的话就可以指定 "none"，结合前面的device比较方便

然后将模型放放在多GPU运行

model = nn.DataParallel(model)

二、代码实例

引入 PyTorch 模块和定义参数

import torch
import torch.nn as nn
from torch.autograd import Variable
from torch.utils.data import Dataset, DataLoader
# Parameters and DataLoaders
input_size = 5
output_size = 2

batch_size = 30
data_size = 100

设备

import os
# 这里可以指定设备的id, 或者 "none"
os.environ["CUDA_VISIBLE_DEVICES"] ="0,1"
DEVICE= torch.device("cuda" if torch.cuda.is_available() else "cpu")
# DEVICE = torch.device("cpu")

生成数据

通过使用Dataset和DataLoader对象

class RandomDataset(Dataset):
    def __init__(self,size,length):
        self.len = length
        self.data = torch.randn(length,size)

    def __getitem__(self, index: int):
        return self.data[index]

    def __len__(self):
        return self.len

rand_dataset = RandomDataset(input_size,data_size)
rand_loader = DataLoader(rand_dataset,batch_size=batch_size,shuffle=True)

简单模型

为了做一个小 demo，我们的模型只是获得一个输入，执行一个线性操作，然后给一个输出。尽管如此，你可以使用 DataParallel 在任何模型(CNN, RNN, Capsule Net 等等.)

我们放置了一个输出声明在模型中来检测输出和输入张量的大小。

class MyModel(nn.Module):
    def __init__(self,input_size,output_size):
        super(MyModel,self).__init__()
        self.fc = nn.Linear(input_size,output_size)

    def forward(self, input):
        output = self.fc(input)
        print("  In Model: input size", input.size(),
              "output size", output.size())
        return output

创建模型并且数据并行处理

这是整个教程的核心。首先我们需要一个模型的实例，然后验证我们是否有多个 GPU。如果我们有多个 GPU，我们可以用 nn.DataParallel 来 包裹 我们的模型。然后我们使用 model.to(device) 把模型放到多 GPU 中。

if torch.cuda.device_count() > 1:
    print("Let's use", torch.cuda.device_count(), "GPUs!")
  # dim = 0 [30, xxx] -> [10, ...], [10, ...], [10, ...] on 3 GPUs
    model = nn.DataParallel(model)

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

＃　前面定义了device，所以可以
model.to(device)

运行程序

for data in rand_loader:
    if torch.cuda.is_available():
        input_var = Variable(data.cuda())
    else:
        input_var = Variable(data)
    # 或者 input_var = data.to(device)

    output = model(input_var)
    print("Outside: input size", input_var.size(),
          "output_size", output.size())

结果

使用1个GPU

使用1个GPU

使用2个GPU

使用2个GPU

使用3个GPU

使用3个GPU

以上就学会了使用GPU加速模型训练，同时知道如何使用多GPU并行化训练。