深度学习开发框架PyTorch(7)-- GPU的使用

PyTorch中使用GPU会复杂一些,但这也意味着对GPU资源更加灵活高效的控制。

一、GPU检查

1、判断是否支持GPU

1
2
3
4
5
6
import torch

if torch.cuda.is_available():
print('支持GPU')
else:
print('不支持GPU')

2、GPU相关api

1
2
3
4
5
6
7
8
9
10
import torch

# GPU数量
print(torch.cuda.device_count())

# GPU型号
print(torch.cuda.get_device_name(0))

# 返回当前设备索引
print(torch.cuda.current_device())

二、GPU的使用

PyTorch的数据结构分为CPU和GPU两个版本,通过调用.cuda方法可以将其转为对应的GPU对象。

1、单机单GPU

构建模型时,将模型转移到GPU上

1
2
3
4
5
6
device = torch.device('cuda:0')
use_gpu = True

if use_gpu and torch.cuda.is_available():
# net = Net().cuda()
net = Net().to(device)

训练模型时,将数据转移到GPU上

1
2
3
for data in trainloader:
inputs, labels = data
inputs, labels = inputs.to(device), labels.to(device)

注意:模型预测结果想要转化为numpy数组需要先把数据转移到CPU

1
2
preditcion = preditcion.cpu()
preditcion = preditcion.data.numpy()

2、单机多GPU(并行)

两种流行的并行方式:模型并行、数据并行。
模型并行指的是一个模型从逻辑上被分成了几个部分(例如,一些层在一部分,其他层在另一部分),然后把它们部署在不同的硬件/设备上。
数据并行指的是,通过位于不同硬件/设备上的同一个网络的多个副本来处理数据的不同批(batch)。不同于模型并行,每个副本可能是整个网络,而不仅仅是一部分。

因为数据并行随着数据的增长可以很好地扩展,所以它更加流行也更加常用。例如batch_size=96和三个GPU,每个GPU就会计算batch_size=32的数据。

1
2
3
4
5
6
7
def forward(self, x):
# 更改前向传播函数,从而使用多块GPU
x = nn.parallel.data_parallel(self.conv1, x, device_ids=None)
x = nn.parallel.data_parallel(self.conv2, x, device_ids=None)
x = x.view(x.size(0), 16*5*5)
x = nn.parallel.data_parallel(self.classifier, x, device_ids=None)
return x

3、多机多GPU(分布式)

DataParallel只能满足一台机器上gpu的通信,而一台机器一般只能装8张卡,对于一些大任务,8张卡就很吃力了,这个时候我们就需要面对多机多卡分布式训练这个问题了。
PyTorch官方给出了应用于分布式训练的api:

  • torch.nn.parallel.DistributedDataParallel
    这个从名字上就能看出来与DataParallel相类似。这个包是实现多机多卡分布训练最核心东西,它可以帮助我们在不同机器的多个模型拷贝之间平均梯度。

  • torch.utils.data.distributed.DistributedSampler
    在多机多卡情况下分布式训练数据的读取也是一个问题,不同的卡读取到的数据应该是不同的。dataparallel的做法是直接将batch切分到不同的卡,这种方法对于多机来说不可取,因为多机之间直接进行数据传输会严重影响效率。于是有了利用sampler确保dataloader只会load到整个数据集的一个特定子集的做法。DistributedSampler就是做这件事的。它为每一个子进程划分出一部分数据集,以避免不同进程之间数据重复。

拓展阅读

0%