Pytorch分布式训练

Pytorch分布式训练DataParallel和DistributedDataParallel详解

本人项目涉及到分布式训练，之前简单介绍了分布式训练神经网络训练——分布式训练（一）
笔者用pytorch比较多，故此记录下pytorch的分布式训练的两种方式，开始学习吧~

1. 入门

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

1. torch.nn.DataParallel ==> 简称 DP，数据并行，PS结构
2. torch.nn.parallel.DistributedDataParallel ==> 简称DDP，数据并行，优于DataParallel，好像仍是PS架构
   其中 DP 只用于单机多卡，DDP 可以用于单机多卡也可用于多机多卡，后者现在也是Pytorch训练的主流用法，DP写法比较简单，但即使在单机多卡情况下也比 DDP 慢。

2. DP（DataParallel ）

之前单机跑模型的时候，发现真真真的是太慢了！像COCO这种大型的数据集，等一个模型跑完，估计都该过年了！！！然后自然而然就想起来去用Dataparallel函数加速并行。
（1）关于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]

可以看出默认把loss计算输出到0号卡上，这会导致0号卡内存占用最多

（2）DP使用

import torch
import torch.nn as nn

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

# 模型放在GPU上
net = net.cuda()
net=nn.DataParallel(net)   #多了此步，默认使用所有的gpu，并输出到0号卡上

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

result = net(inputs)

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

思考：如果想要限制使用的显卡数，怎么办呢？

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号卡)分发。

另一种具体写法：

if torch.cuda.device_count() > 1:#判断是不是有多个GPU
    print("Let's use", torch.cuda.device_count(), "GPUs!")
    # 就这一行
    model = nn.DataParallel(model,device_ids=range(torch.cuda.device_count()))

参数说明：
model：待并行计算的模型
device_ids：并行计算的gpu的id
看起来非常简单对不对，就一行代码就解决了，但是在我使用的时候发现，GPU会出现负载不均衡的问题。为什么呢?就是第一个GPU（12GB）可能占用了10GB，剩余的GPU却只使用了2GB。这就造成的极大的浪费！！！这个问题的原因是，当你在数据并行的时候，模型会复制到每一个gpu作为副本，数据平分，每一块gpu的数据为batch size/gpu数目，你的loss却不是这样的，每次都会在第一个GPU相加计算loss，然后分发到每个gpu，进行梯度计算，再取梯度平均，进行梯度更新。这就造成了第一个GPU的负载远远大于剩余其他的显卡。
怎么解决这个问题呢？？
那么DDP就来了~

3 . DDP（distributedDataparallel）

为了弥补Dataparallel的不足，有torch.nn.parallel.DistributedDataParallel，这也是现在Pytorch分布式训练主推的。
这个函数的主要目的是为了多机多卡加速的，但是单机多卡也是没问题的。相比于之前的Dataparallel，新的函数更加优雅，速度更加快（这一点官方文档里有提到），而且不会出现负载不均衡的问题，唯一的小缺点可能就是配置稍微有点小麻烦。
DDP支持单机多卡和多机多卡，每张卡都有一个进程，这就涉及到进程通信，多进程通信初始化，是DDP使用最复杂的地方。

3.1 DDP的配置

下面主要介绍一下调用过程，我粗略的把他分成三步：
（1)对执行训练的进程进行初始化

#初始化使用nccl后端（这个），当然还有别的后端，可以查看官方文档，介绍的比较清晰
torch.distributed.init_process_group(backend="nccl")
#完整的说明
torch.distributed.init_process_group(backend, init_method='env://', timeout=datetime.timedelta(0, 1800), **kwargs)

参数说明：
backend str/Backend 是通信所用的后端，可以是"ncll" "gloo"或者是一个torch.distributed.Backend类（Backend.GLOO）
init_method str 这个URL指定了如何初始化互相通信的进程，初始化init_method的方法有两种, 一种是使用TCP进行初始化, 另外一种是使用共享文件系统进行初始化。
world_size int 执行训练的所有的进程数
rank int this进程的编号（一般一块gpu一个进程），也是其优先级，rank=0的主机就是主要节点。同时，rank在gpu训练中，又表示了gpu的编号/进程的编号。
timeout timedelta 每个进程执行的超时时间，默认是30分钟，这个参数只适用于gloo后端
group_name str 进程所在group的name
（2）使用DistributedSampler
如果自己写数据流，得根据torch.distributed.get_rank()去shard数据，获取自己应用的一份
如果用Dataset API，则需要在定义Dataloader的时候用DistributedSampler 去shard：

train_dataset = self.dataset(
 self.opt.data_path, train_filenames, self.opt.height, self.opt.width,
 self.opt.frame_ids, 4, is_train=True, img_ext=img_ext)
 
train_sampler = torch.utils.data.distributed.DistributedSampler(train_dataset)

train_loader = DataLoader(
            train_dataset, self.opt.batch_size, False,
 num_workers=self.opt.num_workers, pin_memory=True, drop_last=True,sampler=train_sampler)
 

另外多说一点，在使用DataLoader时，别忘了设置pip_memory=true，为什么呢？就是一个字，快！！
3.分布式训练

model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[arg.local_rank], output_device=arg.local_rank)

到这里你以为就结束了吗？？其实并没有，还剩下一点工作需要做，就是在命令行里加上这样一句话，

python -m torch.distributed.launch --nproc_per_node=NUM_GPUS_YOU_HAVE  yourscript.py

注意事项
1、和DataParalle不同,DistributedDataParallel每个进程上的数据个数都等于batch-size,这导致每一次梯度下降实际用的数据量为batch-size*gpu-num.虽然一般情况下,batch越大计算越快,但是大batch不一定有最好的泛化,对于以及在单GPU或者DataParalle下以及对batch大小精调过的模型,设置batch为之前的GPU个数分之一才能复现之前结果
2、避免写入log文件发生冲突,只在某一个进程里用print,logger.可以通过args.local_rank==0来选择进程
3、计算loss或者accuracy时,要汇总各进程的信息
思考：
1、我觉得DDP的pytorch代码难写，一般使用DP,不如使用pytorch的另一个简单的分布式框架apex，下次有时间我整理下apex，混合精度训练
2、DP可以把loss也并行计算
参考：
超级详细
和nn.DataParallel说再见
Pytorch分布式训练DataParallel和DistributedDataParallel详解
pytorch单机多卡：从DataParallel到distributedDataParallel