不同加速库在 ImageNet 上的使用示例（单机多卡）（Github）：
nn.DataParallel 简单方便的 nn.DataParallel
torch.distributed 使用 torch.distributed 加速并行训练
torch.multiprocessing 使用 torch.multiprocessing 取代启动器
apex 使用 apex 再加速
horovod horovod 的优雅实现
注：分布式 evaluation
测试发现 Apex 的加速效果最好，但与 Horovod/Distributed 差别不大，平时可以直接使用内置的 Distributed。Dataparallel 较慢，不推荐使用。
当代研究生应当掌握的并行训练方法（单机多卡） - 知乎 (zhihu.com)
单GPU代码
## main.py文件
import torch
# 构造模型
model = nn.Linear(10, 10).to(local_rank)
# 前向传播
outputs = model(torch.randn(20, 10).to(rank))
labels = torch.randn(20, 10).to(rank)
loss_fn = nn.MSELoss()
loss_fn(outputs, labels).backward()
# 后向传播
optimizer = optim.SGD(model.parameters(), lr=0.001)
optimizer.step()
## Bash运行
python main.py
## main.py文件
import torch
# 新增：
import torch.distributed as dist
# 新增：从外面得到local_rank参数
import argparse
parser = argparse.ArgumentParser()
parser.add_argument("--local_rank", default=-1)
FLAGS = parser.parse_args()
local_rank = FLAGS.local_rank
# 新增：DDP backend初始化
torch.cuda.set_device(local_rank)
dist.init_process_group(backend='nccl') # nccl是GPU设备上最快、最推荐的后端
# 构造模型
device = torch.device("cuda", local_rank)
model = nn.Linear(10, 10).to(device)
# 新增：构造DDP model
model = DDP(model, device_ids=[local_rank], output_device=local_rank)
# 前向传播
outputs = model(torch.randn(20, 10).to(rank))
labels = torch.randn(20, 10).to(rank)
loss_fn = nn.MSELoss()
loss_fn(outputs, labels).backward()
# 后向传播
optimizer = optim.SGD(model.parameters(), lr=0.001)
optimizer.step()
## Bash运行
# 改变：使用torch.distributed.launch启动DDP模式，
# 其会给main.py一个local_rank的参数。这就是之前需要"新增:从外面得到local_rank参数"的原因
python -m torch.distributed.launch --nproc_per_node 4 main.py
Pytorch DistributedDataParallel简明使用指南 - 知乎 (zhihu.com)
PyTorch 21.单机多卡操作(分布式DataParallel，混合精度，Horovod) - 知乎 (zhihu.com)
