（单机多卡）4种Pytorch并行训练方法

论文地址：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 达到接近线性的可扩展性。

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()