(单机多卡)4种Pytorch并行训练方法

目录

  • 0.查看 GPU 信息
  • 1.nn.DataParallel
  • 2.torch.distributed
  • 3.torch.multiprocessing
  • 4.horovod 实现

论文地址:https://arxiv.org/pdf/2006.15704.pdf
在这篇论文中,来自 Facebook AI 和华沙大学的研究者介绍了 PyTorch 分布式数据并行模型的设计、实现以及评估。

从 v1.5 开始,PyTorch 自身提供了几种加速分布数据并行的技术,包括分桶梯度(bucketing gradients)、通信重叠计算(overlapping computation with communication)以及跳过梯度同步(skipping gradient synchronization)。相关评估结果显示,在配置正确的情况下,PyTorch 分布式数据并行模型可以用 256 个 GPU 达到接近线性的可扩展性。
(单机多卡)4种Pytorch并行训练方法_第1张图片

0.查看 GPU 信息

torch.cuda.is_available()       # 判断 GPU 是否可用
torch.cuda.device_count()       # 判断有多少 GPU
torch.cuda.get_device_name(0)   # 返回 gpu 名字,设备索引默认从 0 开始
torch.cuda.current_device()     # 返回当前设备索引
# 通过字符串
device = torch.device('cpu')
device = torch.device('cuda:1')  # 指定类型及编号。注意,代码不会检查编号是否合法
device = torch.device('cuda')    # 默认为当前设备

1.nn.DataParallel

与PyTorch一起使用GPU非常容易。您可以将模型放在GPU上:

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

Then, you can copy all your tensors to the GPU:

mytensor = my_tensor.to(device)

在多个GPU上执行前向、后向传播是很自然的。但是,PYTORCH默认只使用一个GPU。 你可以轻松地在多个GPU上运行您的操作,方法是让你的模型使用 DataParallel 并行运行:

model = nn.DataParallel(model)

仅使用nn.DataParallel,gpu0和gpu1、gpu0和gpu2、gpu0和gpu3等包含gpu0的组合都是可以的,其余组合不行,报错RuntimeError: module must have its parameters and buffers on device cuda:1 (device_ids[0]) but found one of them on device: cuda:0,因为tensor是默认存在gpu0上的。
https://github.com/tczhangzhi/pytorch-distributed/blob/master/dataparallel.py

2.torch.distributed

分布式PyTorch,主要是Pytorch在v0.2中发布的一个torch.distributed包,我们可以使用import torch.distributed as dist导入使用,分布式Pyrorch允许您在多台机器之间交换Tensors。使用此软件包,您可以通过多台机器和更大的小批量扩展网络训练。

import torch.distributed as dist

dist.init_process_group(backend='tcp',
                        init_method='tcp://[ff15:1e18:5d4c:4cf0:d02d:b659:53ba:b0a7]:23456',
                        world_size=4)

print('Hello from process {} (out of {})!'.format(
        dist.get_rank(), dist.get_world_size()))

3.torch.multiprocessing

torch.multiprocessing 是一个本地 multiprocessing 模块的包装. 它注册了自定义的reducers, 并使用共享内存为不同的进程在同一份数据上提供共享的视图. 一旦 tensor/storage 被移动到共享内存 (见 share_memory_()), 将其发送到任何进程不会造成拷贝开销.
使用时,只需要调用 torch.multiprocessing.spawn,torch.multiprocessing 就会帮助我们自动创建进程。如下面的代码所示,spawn 开启了 nprocs=4 个线程,每个线程执行 main_worker 并向其中传入 local_rank(当前进程 index)和 args(即 4 和 myargs)作为参数:

import torch.multiprocessing as mp

mp.spawn(main_worker, nprocs=4, args=(4, myargs))

4.horovod 实现

https://github.com/tczhangzhi/pytorch-distributed/blob/master/horovod_distributed.py
与 torch.distributed.launch 相似,我们只需要编写一份代码,horovodrun 启动器就会自动将其分配给n个进程,分别在n个 GPU 上运行。在执行过程中,启动器会将当前进程的(其实就是 GPU的)index 注入 hvd,我们可以这样获得当前进程的 index:

# main.py
import torch
import horovod.torch as hvd

hvd.init()
torch.cuda.set_device(hvd.local_rank())

train_dataset = ...
train_sampler = torch.utils.data.distributed.DistributedSampler(
    train_dataset, num_replicas=hvd.size(), rank=hvd.rank())

train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=..., sampler=train_sampler)

model = ...
model.cuda()

optimizer = optim.SGD(model.parameters())

optimizer = hvd.DistributedOptimizer(optimizer, named_parameters=model.named_parameters())
hvd.broadcast_parameters(model.state_dict(), root_rank=0)

for epoch in range(100):
   for batch_idx, (data, target) in enumerate(train_loader):
       images = images.cuda(non_blocking=True)
       target = target.cuda(non_blocking=True)
       ...
       output = model(images)
       loss = criterion(output, target)
       ...
       optimizer.zero_grad()
       loss.backward()
       optimizer.step()

(单机多卡)4种Pytorch并行训练方法_第2张图片
(单机多卡)4种Pytorch并行训练方法_第3张图片

你可能感兴趣的:(Python相关,算法,人工智能,算法,人工智能,启发式算法)