从 DP 到 DDP 到 apex

1、从 DataParallel 到 DistributedDataParallel，最大的改变是启动方式：

从 python a.py 到 python -m torch.distributed.launch --nproc_per_node=4 --use_env a.py

--nproc_per_node 表示单机使用4块GPU， --use_env 表示不给 a.py 增加参数--local_rank。

torch.distributed.launch 的作用就是用多进程的方式启动 a.py，并增加一些环境变量，比如 ‘RANK’, 'LOCAL_RANK', 'WORLD_SIZE' 等。

import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP

#假设2台机器，每台4卡
_rank = int(os.environ['RANK'])          #全局索引，从0到7
_gpu_id = int(os.environ['LOCAL_RANK'])  #2台机器都是从0到3
_gpu_num = int(os.environ['WORLD_SIZE']) #总进程数=2*4

print('pid=%d, _rank=%d'%(os.getpid(), _rank))
print('pid=%d, _gpu_id=%d'%(os.getpid(), _gpu_id))
print('pid=%d, _gpu_num=%d'%(os.getpid(), _gpu_num))

torch.cuda.set_device(_gpu_i