torch分布式训练DataParallel和DistributedDataParallel

Pytorch 分布式训练主要有两种方式：

torch.nn.DataParallel ==> 简称 DP
torch.nn.parallel.DistributedDataParallel ==> 简称DDP

其中 DP 只用于单机多卡，DDP 可以用于单机多卡也可用于多机多卡，后者现在也是Pytorch训练的主流用法，DP写法比较简单，但即使在单机多卡情况下也比 DDP 慢。

1 DP

import torch
import torch.nn as nn

# 构造模型
net = model(imput_size, output_size)

# 模型放在GPU上
net = net.cuda()
net=nn.DataParallel(net)

# 数据放在GPU上
inputs, labels = inputs.cuda(), labels.cuda()

result = net(inputs)

# 其他和正常模型训练无差别

• 示例
  

2 DDP

import os
import re

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

# 1. 获取环境信息
rank = int(os.environ['SLURM_PROCID'])
world_size = int(os.environ['SLURM_NTASKS'])
local_rank = int(os.environ['SLURM_LOCALID'])
node_list = str(os.environ['SLURM_NODELIST'])       

# 对ip进行操作
node_parts = re.findall('[0-9]+', node_list)
host_ip = '{}.{}.{}.{}'.format(node_parts[1], node_parts[2], node_parts[3], node_parts[4])

 # 注意端口一定要没有被使用
port = "23456"                                         

 # 使用TCP初始化方法
init_method = 'tcp://{}:{}'.format(host_ip, port)      

# 多进程初始化,初始化通信环境
dist.init_process_group("nccl", init_method=init_method,
                        world_size=world_size, rank=rank) 

# 指定每个节点上的device
torch.cuda.set_device(local_rank)
                     
model = model.cuda()

# 当前模型所在local_rank
model = DDP(model, device_ids=[local_rank])             # 指定当前卡上的GPU号

input = input.cuda()
output = model(input)

# 此后训练流程与普通模型无异