Pytorch分布式训练DataParallel和DistributedDataParallel

原创如下：

https://www.codenong.com/cs105605994/

下面的记录只是方便自己翻阅，建议去原博客观看

---

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

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

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

一、DP(DataParallel)

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)
# 其他和正常模型训练无差别

关于Dataparallel, 摘取主要源码:

class DataParallel(Module):
    def __init__(self, module, device_ids=None, output_device=None, dim=0):
        super(DataParallel, self).__init__()
       
        # 如果没有GPU可用，直接返回
        if not torch.cuda.is_available():
            self.module = module
            self.device_ids = []
            return

        # 如果有GPU，但没有指定的话，device_ids为所有可用GPU
        if device_ids is None:
            device_ids = list(range(torch.cuda.device_count()))
           
        # 默认输出在0号卡上
        if output_device is None:
            output_device = device_ids[0]

总结
如果不设定好要使用的device_ids的话, 程序会自动找到这个机器上面可以用的所有的显卡用于训练。
如果想要限制使用的显卡数，怎么办呢？
在代码最前面使用：

os.environ['CUDA_VISIBLE_DEVICES'] == '0,5'
# 限制代码能看到的GPU个数，这里表示指定只使用实际的0号和5号GPU
# 注意：这里的赋值必须是字符串，list会报错

# 这时候device_count = 2
device_ids = range(torch.cuda.device_count())

# device_ids = [0,1] 这里的0就是上述指定的'0'号卡，1对应'5'号卡。
net = nn.DataParallel(net，device_ids)

# !!!模型和数据都由主gpu(0号卡)分发。

值得注意的是，在使用os.environ[‘CUDA_VISIBLE_DEVICES’]对可以使用的显卡进行限定之后, 显卡的实际编号和程序看到的编号应该是不一样的。

例如上面我们设定的是os.environ[‘CUDA_VISIBLE_DEVICES’]=“0,5”, 但是程序看到的显卡编号应该被改成了’0,1’。

也就是说程序所使用的显卡编号实际上是经过了一次映射之后才会映射到真正的显卡编号上面的, 例如这里的程序看到的1对应实际的5。

但是Dataparallel会带来显存的使用不平衡，具体分析见参考链接[2]，而且碰到大的任务，时间和能力上都很受限。

二、DDP (DistributedDataParallel)

为了弥补Dataparallel的不足，有torch.nn.parallel.DistributedDataParallel，这也是现在Pytorch分布式训练主推的。

DDP支持单机多卡和多机多卡，每张卡都有一个进程，这就涉及到进程通信，多进程通信初始化，是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)

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