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)
# 此后训练流程与普通模型无异
