從 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_id)
dist.init_process_group(backend='nccl')
dist.barrier()

...

#train_dataset 是一個 Dataset 的執行個體
train_sampler = torch.utils.data.distributed.DistributedSampler(train_dataset)
data_train = DataLoader(train_dataset,                                                
                        batch_size=int(batchsize/_gpu_num),
                        shuffle=False, #True, 在 DistributedSampler 裡做過shuffle了
                        num_workers=4,
                        sampler=self.train_sampler)

#model = DataParallel(model)
model = DDP(model, device_ids=[_gpu_id], output_device=_gpu_id)

...

#在每個epoch訓練的開始
train_sampler.set_epoch(epoch_idx)

#在所有寫日志和儲存模型的地方加上對 _rank 的判斷
if (_rank == 0):
    writelog or savemodel

           

2、從DistributedDataParallel 到 apex，最大的坑是 GCC 需要大于 4.9。（apex 已放棄，pytorch1.6自帶的amp好用）

from apex import amp
#from torch.nn.parallel import DistributedDataParallel as DDP
from apex.parallel import DistributedDataParallel as DDP

...

#2個模型，1個優化器, 先做initialize，再做DDP
list_model = [model_a, model_b]
list_model, optimizer = amp.initialize(list_model, optimizer, opt_level="O1")
model_a = DDP(model_a)
model_b = DDP(model_b)

...

#loss.backward()
with amp.scale_loss(loss, optimizer) as scaled_loss:
    scaled_loss.backward()

           

參考：

1、[原創][深度][PyTorch] DDP系列第一篇：入門教程

2、scaled_loss.backward Segmentation fault (core dumped)

3、Pytorch中的Distributed Data Parallel與混合精度訓練（Apex）