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系列第一篇:入门教程