torch.distributed多卡/多GPU/分布式DPP(二)——torch.distributed.all_reduce(reduce_mean)＆barrier＆控制进程执行顺序＆随机数种子

torch.distributed.all_reduce

实现的就是Ring-Reduce版本

跨卡汇总

torch.distributed.all_reduce(tensor, op=<ReduceOp.SUM: 0>, group=None, async_op=False)

举例

import torch
import argparse
import torch.distributed as dist
from torch.distributed import ReduceOp
# 新增1：从外面得到local_rank参数，在调用DDP的时候，其会自动给出这个参数
parser = argparse.ArgumentParser()
parser.add_argument("--local_rank", default=-1, type=int)
args = parser.parse_args()

# 新增2：DDP backend初始化
#   a.根据local_rank来设定当前使用哪块GPU
torch.cuda.set_device(args.local_rank)
#   b.初始化DDP, nccl是GPU设备上最快、最推荐的后端.如果是CPU模型运行，需要选择其他后端。
dist.init_process_group(backend='nccl')
world_size = torch.distributed.get_world_size()
print('全局进程数: ', world_size)
 
# 新增3：定义并把模型放置到单独的GPU上
device = torch.device("cuda", args.local_rank)

tensor = (torch.arange(world_size, dtype=torch.int64) + 1 + world_size * args.local_rank).to(device)
print(tensor)

dist.all_reduce(tensor, op=ReduceOp.SUM)
print(tensor)

python -m torch.distributed.launch \
       --nproc_per_node 2 \
       main.py

4 = 3+1

6 = 4+2

如果是求均值的话(reduce_mean)

dist.all_reduce(value)
if average:
    value /= world_size

torch.distributed.barrier DDP各个进程同步 

torch.distributed提供了一个barrier()的接口，利用它我们可以同步各个DDP中的各个进程！当使用barrier函数时，DDP进程会在函数的位置进行等待，知道所有的进程都跑到了 barrier函数的位置，它们才会再次向下执行。

torch.distributed.barrier(group=None, async_op=False, device_ids=None)

用torch.distributed.barrier控制DDP不同进程的执行顺序

一般情况下，各个进程是各自执行的，速度有快有慢，只有在gradient all-reduce的时候，快的进程才会等一下慢的进程，也就是进行同步

在某个进程中执行A操作，其他进程等待其执行完成后再执行B操作

if rank == 0:
    do_A()
    torch.distributed.barrier()
else:
    torch.distributed.barrier()
    do_B()

在加载数据前，如果数据集不存在，我们要下载数据集

1. 我们只需要在唯一一个进程中开启一次下载
2. 我们需要让其他进程等待其下载完成，再去加载数据

怎么解决这个问题呢？torch.distributed提供了一个barrier()的接口，利用它我们可以同步各个DDP中的各个进程！当使用barrier函数时，DDP进程会在函数的位置进行等待，知道所有的进程都跑到了 barrier函数的位置，它们才会再次向下执行。

在某个进程中优先执行A操作，其他进程等待其执行完成后再执行A操作

这个值得深入讲一下，因为这个是非常普遍的需求。利用contextlib.contextmanager，我们可以把这个逻辑给优雅地包装起来！

from contextlib import contextmanager

@contextmanager
def torch_distributed_zero_first(rank: int):
    """Decorator to make all processes in distributed training wait for each local_master to do something.
    """
    if rank not in [-1, 0]:
        torch.distributed.barrier()
    # 这里的用法其实就是协程的一种哦。
    yield
    if rank == 0:
        torch.distributed.barrier()

然后我们就可以这样骚操作：

with torch_distributed_zero_first(rank):
    if not check_if_dataset_exist():
        download_dataset()
    load_dataset()

给不同的进程分配不同的、固定的随机数种子

为保证实验的可复现性，一般我们会在代码在开头声明一个固定的随机数种子，从而使得同一个配置下的实验，无论启动多少次，都会拿到同样的结果。

import random
import numpy as np
import torch

def init_seeds(seed=0, cuda_deterministic=True):
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    # Speed-reproducibility tradeoff https://pytorch.org/docs/stable/notes/randomness.html
    if cuda_deterministic:  # slower, more reproducible
        cudnn.deterministic = True
        cudnn.benchmark = False
    else:  # faster, less reproducible
        cudnn.deterministic = False
        cudnn.benchmark = True
        

def main():
    # 一般都直接用0作为固定的随机数种子。
    init_seeds(0)

但是在DDP训练中，如果还是像以前一样，使用0作为随机数种子，不做修改，就会造成以下后果：

1. DDP的N个进程都使用同一个随机数种子
2. 在生成数据时，如果我们使用了一些随机过程的数据扩充方法，那么，各个进程生成的数据会带有一定的同态性。
   1. 比如说，YOLOv5会使用mosaic数据增强（从数据集中随机采样3张图像与当前的拼在一起，组成一张里面有4张小图的大图）。这样，因为各卡使用了相同的随机数种子，你会发现，各卡生成的图像中，除了原本的那张小图，其他三张小图都是一模一样的！
3. 同态性的数据，降低了训练数据的质量，也就降低了训练效率！最终得到的模型性能，很有可能是比原来更低的。

所以，我们需要给不同的进程分配不同的、固定的随机数种子：

def main():
    rank = torch.distributed.get_rank()
    # 问题完美解决！
    init_seeds(1 + rank)