看懂代码中的pytorch并行计算

torch.multiprocessing.spawn——执行多进程并行任务

1. torch.multiprocessing模块

torch.multiprocessing是pytorch中用于多进程并行计算的模块

2. torch.multiprocessing.spawn方法

用于在多个进程中，并行地执行指定的函数

和torch.distributed.lauch/ torchrun不同，这两个是命令行工具，用于并行执行指定的Python脚本

import torch.multiprocessing as mp

mp.spawn(fn, args=(), nprocs=1, join=True, deamon=False, start_method='spawn')

参数：

• fn：目标函数
• args： 传递给目标函数地参数，是一个元组
• nprocs：要启动的进程数量
• join：是否等待所有进程完成后再返回
• daemon：如果为True，则子进程会作为守护进程运行
• start_method：进程启动的方法，可以是‘spawn’,'fork',或'forkserver', 默认是‘spawn’，因为在多GPU场景中更安全

示例：

import torch.multiprocessing as mp

def example(rank, world_size):
    print(f"I am process {rank} of {world_size}")

if __name__ == "__main__":
    world_size = 4
    mp.spawn(example, args=(world_size,), nprocs = world_size, join=True)

I am process 0 of 4

I am process 1 of 4

I am process 2 of 4

I am process 3 of 4

• 函数的第一个参数rank表示进程的索引。这个rank是mp.spawn函数自动传递的（从0到world_size-1），不需要手动传递

3. Moco源代码中的使用

在Moco的源代码中也用到了这个方法，具体的使用情况如下：

首先使用torch.cuda.device_count()自动获取设备可用的GPU数量，将它赋值给nprocs，同时也是main_worker函数的一个参数。

在main_worker函数中，第一个参数gpu就是spawn方法自动传递的，从0到ngpus_per_node-1.

ngpus_per_node = torch.cuda.device_count()
mp.spawn(main_worker, nprocs=ngpus_per_node, args=(ngpus_per_node, args))
def main_worker(gpu, ngpus_per_node, args):
    args.gpu = gpu

    if args.gpu is not None:
        print("Use GPU: {} for training".format(args.gpu))
    # suppress printing if not master
    if args.multiprocessing_distributed and args.gpu != 0:
    # 这一块的作用是只在第0个进程的时候打印我们想要打印的东西，当args.gpu !=0 时，
    # 后续的print函数都被重置为pass，也就是不输出任何东西

        def print_pass(*args):
            pass

        builtins.print = print_pass
     ……………………

初始化进程组 torch.distributed.init_process_group

使用DDP之前，首先要正确设置进程组(在函数fn里面进行初始化)。

torch.distributed.init_process_group(
    backend,
    init_method=None,
    timeout=datetime.timedelta(seconds=1800),
    world_size=-1,
    rank=-1,
    store=None,
    group_name='default',
    **kwargs
)

参数：

• backend：指定分布式后端的名称，例如 ‘nccl’、‘gloo’ 或 ‘mpi’。
• init_method：初始化方法的 URL 或文件路径。默认为 None，表示使用默认的初始化方法。
• timeout：初始化过程的超时时间，默认为 1800 秒。
• world_size：参与分布式训练的总进程数。默认为 -1，表示从环境变量中自动获取。
• rank：当前进程的排名。默认为 -1，表示从环境变量中自动获取。
• store：用于存储进程组信息的存储对象。默认为 None，表示使用默认存储。
• group_name：进程组的名称，默认为 ‘default’。
• **kwargs：其他可选参数，根据不同的分布式后端而定。

Moco代码中的使用：

if args.distributed:
    if args.dist_url == "env://" and args.rank == -1:
        args.rank = int(os.environ["RANK"])
    if args.multiprocessing_distributed:
        # For multiprocessing distributed training, rank needs to be the
        # global rank among all the processes
        args.rank = args.rank * ngpus_per_node + gpu
        # args.rank原本为0，因此args.rank=gpu，也就是当前的进程编号
    dist.init_process_group(
        backend=args.dist_backend,
        init_method=args.dist_url,
        world_size=args.world_size,
        rank=args.rank
        # args.world_size为-1
    )

数据并行DDP——torch.nn.parallel.DistributedDataParallel

torch.nn.parallel.DistributedDataParallel(model, 
                                          device_ids=None,
                                          output_device=None
                                          )

参数：

• model：要并行训练的模型
• device_ids：用于指定模型将被复制到哪些GPU上
• output_device：用于指定要将输出转移到哪个GPU上。默认值为None，表示将输出保留在各自的GPU上，不进行转移，一般使用默认值即可.

Moco代码中对应的部分：

1. 第一步先对数据进行采样

   if args.distributed:
       train_sampler = torch.utils.data.distributed.
           DistributedSampler(train_dataset)
   
   
   train_loader = torch.utils.data.DataLoader(
           train_dataset,
           batch_size=args.batch_size,
           shuffle=(train_sampler is None), # 并行计算时shuffle=False
           num_workers=args.workers,
           pin_memory=True,
           sampler=train_sampler,
           drop_last=True,
       )

2. 第二步获得local_rank(在这里就是args.gpu，也就是当前进程编号)
3. 使用torch.cuda.set_device(loacal_rank)，将进程和GPU进行绑定，即指定当前进程应该使用哪个GPU.

   # (在main_worker函数中)
   torch.cuda.set_device(args.gpu)
   model.cuda(args.gpu)
   model = torch.nn.parallel.DistributedDataParallel(
                   model, device_ids=[args.gpu]
               )