pytorch——关于GPU的一些知识

一、前言

最近在学习pytorch框架，记录一些涉及到的知识点，方便后续查找和学习。

二、主要内容

内容可能有些散，初次学习，可能把握不好知识之间的连贯性和整体性，后续适当调整。

关于CUDA的一些函数接口：

torch.cuda.is_available()  #查看系统GPU是否可以使用，经常用来判断是否装好gpu版的pytorch

torch.cuda.current_device()  #返回当前设备序号

torch.cuda.get_device_name(0) #返回第0号设备的name

torch.cuda.device_count() #返回可使用的gpu的数量

torch.cuda.memory_allocated(device=‘cuda:0’)  #返回0号设备的当前GPU显存使用量(以字节为单位）

torch.device

作为Tensor 的一种属性，其包含了两种设备类型，cpu和gpu，通常通过两种方式来进行创建：

1.通过字符串创建
eg: 
torch.device('cuda:0')
torch.device('cpu')
torch.device('cuda')  # 当前cuda设备

2.通过字符串加设备编号创建
eg:
torch.device('cuda', 0)
torch.device('cpu', 0)

常见的几种创建gpu设备上的Tensor对象:
torch.randn((2,3),device = torch.device('cuda:0'))
torch.randn((2,3),device = 'cuda:0')
torch.randn((2,3),device = 0)  #历史遗留做法，仅支持gpu

.to()

进行设备转化，也是常用的设置gpu的方式。

通常可以这样调用:

to(device=None, dtype=None, non_blocking=False)
#第一个可以设置当前的设备，比如device=torch.device('cuda:0')
#第二个就是数据类型了，如torch.float,torch.int,torch.double
#第三个参数，若设置为True, 并且此对象的资源存储在固定内存上(pinned memory)，那么此cuda()函数产生的复制将与host端的原storage对象保持同步。否则此参数不起作用

使用指定的GPU:

PyTorch默认使用从0开始的GPU,常有以下两种方式来指定特定的GPU

1.CUDA_VISIBLE_DEVICES
终端设置:  CUDA_VISIBLE_DEVICES=1,2  python train.py    (举个例子)
代码里设置:
			import os
			os.environ["CUDA_VISIBLE_DEVICES"] = '1,2'
2.torch.cuda.set_device()
代码里设置:
			import torch
			torch.cuda.set_device(id)

不过官方建议使用CUDA_VISIBLE_DEVICES，不建议使用 set_device 函数。

----- 分割线----

了解了以上的一些gpu的操作后，我们来说一些实际训练时的方法。

单卡训练:

这个比较简单的一种操作方式。例子如下：

#模型置于gpu
device = torch.device("cuda: 0")
model.to(device)   # or   model.cuda(device),数据处理类似可这样

#数据置于gpu
mytensor = my_tensor.to(device)  #这里注意，mytensor 是my_tensor的一个gpu上的一个副本，而不是重写了my_tensor

多卡训练:

为了提高训练速度，常常一个机器有多张gpu，这时候便可以进行并行训练，提高效率。而并行训练又可分为数据并行处理和模型并行处理。
数据并行处理指的是使用同一个模型，将数据均匀分到多种gpu上，进行训练；
模型并行处理指的是多张gpu训练模型的不同部分，使用同一批数据。
下面将就这个两个方法进行记录。

数据并行处理:

当我们的数据量太大时，可以考虑数据并行处理方法。

API:
class torch.nn.DataParallel(module, device_ids=None, output_device=None, dim=0)

# module:即所用到的模型
# device_ids:进行数据并行处理的设备编号，默认使用设备的所有gpu。若指定gpu，常用编号组成的int列表
# output_device:输出数据的位置，默认为第一个编号的gpu

使用方法：

import torch.nn as nn

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

if torch.cuda.device_count() > 1: 
	nn.DataParallel(model)  #也可根据已有gpu进行指定，eg:nn.DataParallel(model,device_ids = [0,1,2])
	
model.to(device)

这里稍微解释下：上述处理模型的操作，数据处理的操作和单卡一样，放于设定的device上即可，可能有人会问，这不是和单卡训练一样嘛。其实是不一样的，nn.DataParallel()的设定，使得模型进行处理数据时，会将model复制到选定的gpu上，然后将批量的数据进行均分，分配到选定gpu上，当然，这就隐含一个条件，批量的数据必须是大于选定的gpu数量的，要不怎么分？然后进行计算，最后将结果返回选定的第一个gpu上，这里注意一点，所选的gpu必须要包含device所设定gpu，可以理解为主gpu。
比较详细的例子见：pytorch官网关于数据并行处理的例子

模型并行处理:

这种方法，自然对应的是模型太大的情况。当模型较大时，将模型的不同层放于不同的gpu上，进行并行的处理。一个简单的例子帮助理解，见下：

import torch
import torch.nn as nn
import torch.optim as optim

class ToyModel(nn.Module):
    def __init__(self):
        super(ToyModel, self).__init__()
        self.net1 = torch.nn.Linear(10, 10).to('cuda:0')
        self.relu = torch.nn.ReLU()
        self.net2 = torch.nn.Linear(10, 5).to('cuda:1')

    def forward(self, x):
        x = self.relu(self.net1(x.to('cuda:0')))
        return self.net2(x.to('cuda:1'))

model = ToyModel()
loss_fn = nn.MSELoss()
optimizer = optim.SGD(model.parameters(), lr=0.001)

optimizer.zero_grad()
outputs = model(torch.randn(20, 10))
labels = torch.randn(20, 5).to('cuda:1')
loss_fn(outputs, labels).backward()
optimizer.step()

# 其中，backward()和torch.optim将自动处理梯度，只需确保调用损失函数时标签与输出在同一设备上。

然而，上述方法虽解决了模型太大带来的问题，但其花费时间是大于单gpu的。原因是在任何时间点，两个GPU中只有一个在工作，而另一个在那儿什么也没做。
如下图，是pytorch官网实现的reset50ModelParallelResNet50和torchvision.models.reset50()的消耗时间对比图：（具体例子可见pytorch官网）

后又用流水线输入加速：
eg：以下例子来自于pytorch官网

import torchvision.models as models
import matplotlib.pyplot as plt
plt.switch_backend('Agg')
import numpy as np
import timeit

num_batches = 3
batch_size = 120
image_w = 128
image_h = 128

from torchvision.models.resnet import ResNet, Bottleneck

num_classes = 1000


class ModelParallelResNet50(ResNet):
    def __init__(self, *args, **kwargs):
        super(ModelParallelResNet50, self).__init__(
            Bottleneck, [3, 4, 6, 3], num_classes=num_classes, *args, **kwargs)

        self.seq1 = nn.Sequential(
            self.conv1,
            self.bn1,
            self.relu,
            self.maxpool,

            self.layer1,
            self.layer2
        ).to('cuda:0')

        self.seq2 = nn.Sequential(
            self.layer3,
            self.layer4,
            self.avgpool,
        ).to('cuda:1')

        self.fc.to('cuda:1')

    def forward(self, x):
        x = self.seq2(self.seq1(x).to('cuda:1'))
        return self.fc(x.view(x.size(0), -1))

def train(model):
    model.train(True)
    loss_fn = nn.MSELoss()
    optimizer = optim.SGD(model.parameters(), lr=0.001)

    one_hot_indices = torch.LongTensor(batch_size) \
                           .random_(0, num_classes) \
                           .view(batch_size, 1)

    for _ in range(num_batches):
        # generate random inputs and labels
        inputs = torch.randn(batch_size, 3, image_w, image_h)
        labels = torch.zeros(batch_size, num_classes) \
                      .scatter_(1, one_hot_indices, 1)

        # run forward pass
        optimizer.zero_grad()
        outputs = model(inputs.to('cuda:0'))

        # run backward pass
        labels = labels.to(outputs.device)
        loss_fn(outputs, labels).backward()
        optimizer.step()

class PipelineParallelResNet50(ModelParallelResNet50):
    def __init__(self, split_size=20, *args, **kwargs):
        super(PipelineParallelResNet50, self).__init__(*args, **kwargs)
        self.split_size = split_size

    def forward(self, x):
        splits = iter(x.split(self.split_size, dim=0))
        s_next = next(splits)
        s_prev = self.seq1(s_next).to('cuda:1')
        ret = []

        for s_next in splits:
            # A. s_prev runs on cuda:1
            s_prev = self.seq2(s_prev)
            ret.append(self.fc(s_prev.view(s_prev.size(0), -1)))

            # B. s_next runs on cuda:0, which can run concurrently with A
            s_prev = self.seq1(s_next).to('cuda:1')

        s_prev = self.seq2(s_prev)
        ret.append(self.fc(s_prev.view(s_prev.size(0), -1)))

        return torch.cat(ret)

num_repeat = 10

# model parallel
stmt = "train(model)"
setup = "model = ModelParallelResNet50()"

mp_run_times = timeit.repeat(
    stmt, setup, number=1, repeat=num_repeat, globals=globals())
mp_mean, mp_std = np.mean(mp_run_times), np.std(mp_run_times)

# Single GPU
setup = "import torchvision.models as models;" + \
        "model = models.resnet50(num_classes=num_classes).to('cuda:0')"
rn_run_times = timeit.repeat(
    stmt, setup, number=1, repeat=num_repeat, globals=globals())
rn_mean, rn_std = np.mean(rn_run_times), np.std(rn_run_times)

# Pipelining model parallel 
setup = "model = PipelineParallelResNet50()"
pp_run_times = timeit.repeat(
    stmt, setup, number=1, repeat=num_repeat, globals=globals())
pp_mean, pp_std = np.mean(pp_run_times), np.std(pp_run_times)

#绘制
plot([mp_mean, rn_mean, pp_mean],
     [mp_std, rn_std, pp_std],
     ['Model Parallel', 'Single GPU', 'Pipelining Model Parallel'],
     'mp_vs_rn_vs_pp.png')
     

最后时间对比图：

当然，不同的split 会有不同的效果，这个在官网也说了可以寻找最佳的split 设置。

三、结束语

以上就是一些gpu的操作知识，关于后面说到的数据并行处理和模型并行处理，只用过数据并行处理，模型并行处理未使用过，所以其中的一些细节可能把握不是很好。而且后面还有分布式数据并行，DistributedDataParallel（DDP），这个可以与模型并行处理一起使用，DataParallel是不可以的，这样更加大大加快模型训练速度，但限于时间，还未学习这一块内容，后续有时间在进行完善一下。

参考链接1：PyTorch中使用指定的GPU
参考链接2：Python：timeit模块详解
参考链接3：pytorch中文翻译版官网