第二十四章解读Pytorch中多GPU并行计算教程(工具)

本文链接：https://blog.csdn.net/weixin_44302770/article/details/134653104

在GPU上运行pytorch程序（指定单/多显卡)
torch.cuda常用指令
 使用CUDA_VISIBLE_DEVICES设置显卡

本教程所涉及的代码github上下载：
https://github.com/WZMIAOMIAO/deep-learning-for-image-processing
在pytorch_classification模块下的train_multi_GPU文件夹中。

常见多GPU使用方法

在训练模型中，为了加速训练过程，往往会使用多块GPU设备进行并行训练（甚至多机多卡的情况）。如下图所示，常见的多GPU的使用方法有以下两种（但不局限于以下方法）：

model parallel，当模型很大，单块GPU的显存不足以放下整个模型时，通常会将模型分成多个部分，每个部分放到不同的GUP设备中（下图左侧），这样就能将原本跑不了的模型利用多块GPU跑起来。但这种情况，一般不能加速模型的训练。
data parallel，当模型不是很大可以放入单块GPU时，可以将模型复制到多块GPU上，进行并行加速训练（下图右侧）。这种情况更常见，本文也是以data parallel来进行讲解。

常见多GPU使用方法

下图展示了使用多块GPU并行加速的训练时间对比。测试环境，Pytorch1.7，CUDA10.1，Model: ResNet34，DataSet: flower_photos，BatchSize: 16，GPU: Tesla V100。通过左侧的柱状图可以看出，使用多GPU的加速并不是简单的线性倍增关系，因为多GPU并行训练时会涉及多GPU之间的通信。
多GPU加速对比

多GPU并行训练过程中需要注意的事项

以下说的注意事项，虽然Pytorch框架已为我们实现了，但是我们需要知道有这些工作。

数据如何分配至各设备当中。使用多GPU并行训练时，通常每个GPU只负责整个数据集中的某一部分。
误差梯度如何在不同的设备之间进行通信。每次多块GPU设备正向传播一批数据后，在误差反向传播时每个GPU设备都会计算出针对各输入数据在各参数的误差梯度（Gradient如下图右侧所示），此时不要急着去更新各参数，而是先去对各设备上各参数的误差梯度求均值（可理解为融合各设备上学习的知识），然后再去更新各设备上参数。
BatchNormalization如何在不同设备间同步。关于BN理论知识不在本文介绍范围内，如果不了解的可以查看我之前写的一篇文章，Batch Normalization理论详解。如果不考虑多设备之间的BN通信的话，每个设备只去计算每个BN层针对该设备输入数据的均值和方差。假如每个设备的batch_size为2，则每个BN层计算的均值和方差只是针对2个样本的。之前在讲BN理论时有说过，一般batch_size设置越大效果越好，那么如果我们在计算BN层的均值和方差时能够同步多块GPU上的统计信息，那batch_size不就相当于倍增了？确实如此，在Pytorch中也有提供具有同步BN的方法SyncBatchNorm。当GPU显存有限，每个设备上的batch_size设置很小时，通过使用具有同步功能的BN层时是能够提升模型最终的mAP的，但如果每个设备上的batch_size设置的已经很大了，那么个人感觉同步的BN就没太大作用了。注意：如果使用具有同步功能的BN，会降低模型的训练速度，因为在每个BN层处都需要去同步参数，所以会更耗时。

下图展示了使用单GPU训练，多GPU训练（使用SyncBatchNorm和不使用SyncBatchNorm）的训练曲线。通过以下曲线可以看出，使用多GPU训练（不使用SyncBatchNorm）和单GPU的训练结果是差不多的（但多GPU训练更快）。但使用了SyncBatchNorm比不使用SyncBatchNorm能达到的最好mAP要高一点。
mAP对比

Pytorch中提供的两种多GPU训练方法

在Pytorch当中，提供了两种多GPU的训练方法，一种是DataParallel一种是DistributedDataParallel，前者是官方较早提供的一种方法，后者是现在官方比较推荐的一种方法。本文也主要是讲DistributedDataParallel。下图是我从官方教程中截取的一段对比这两种方法的文献。
首先DataParallel是单进程多线程的方法，并且仅能工作在单机多卡的情况。而DistributedDataParallel方法是多进程，多线程的，并且适用与单机多卡和多机多卡的情况。即使在在单机多卡的情况下DistributedDataParallell也比DataParallel的速度更快。
本文只介绍单机多卡的情况。

pytorch中提供的多GPU训练方法

Pytorch中多GPU常用启动方式

在Pytorch中使用多GPU的常用启动方式一种是torch.distributed.launch一种是torch.multiprocessing模块。这两种方式各有各的好处，在我使用过程中，感觉torch.distributed.launch启动方式更方便，而且我看官方提供的多GPU训练FasterRCNN源码就是使用的torch.distributed.launch方法，所以我个人也比较喜欢这个方法。但在官方的教程中主要还是使用的torch.multiprocessing方法，官方说这种方法具有更好的控制和灵活性。在自己使用体验过程中确实和官方说的一样。
这里提醒下要使用torch.distributed.launch启动方式的小伙伴。训练过程中如果你强行终止的程序，在开启下次训练前建议你通过nvidia-smi指令看下你GPU的显存是否全部释放了，如果没有全部释放，需要手动杀下进程。在我使用过程中发现强行终止程序有小概率出现进程假死的情况，占用的GPU的资源并没有及时释放，如果在下次训练前没有及时释放，会影响你的训练，或者直接提示通信端口被占用，无法启动的情况。
在我提供的代码中，分别提供了对应这两种方式的训练脚本。torch.distributed.launch对应的是train_multi_gpu_using_launch.py脚本，torch.multiprocessing对应的是train_multi_gpu_using_spawn.py脚本。在这里插入图片描述

train_multi_gpu_using_launch.py脚本讲解

该代码是在之前所讲的知识基础上进行扩展的，其中涉及resnet模型的搭建以及自定义数据集，这里就不在赘述，如果需要了解的可以看下我之前的视频：

这里只是针对其中我个人觉得比较重要的地方说一下，如果想看整个代码的详细讲解，可以去看下本文开头提供的视频链接。

首先说下init_distributed_mode函数，该函数是用来初始化各进程的：

def init_distributed_mode(args):
    if 'RANK' in os.environ and 'WORLD_SIZE' in os.environ:
        args.rank = int(os.environ["RANK"])
        args.world_size = int(os.environ['WORLD_SIZE'])
        args.gpu = int(os.environ['LOCAL_RANK'])
    else:
        print('Not using distributed mode')
        args.distributed = False
        return

    args.distributed = True

    torch.cuda.set_device(args.gpu)
    args.dist_backend = 'nccl'  # 通信后端，nvidia GPU推荐使用NCCL
    print('| distributed init (rank {}): {}'.format(
        args.rank, args.dist_url), flush=True)
    dist.init_process_group(backend=args.dist_backend, init_method=args.dist_url,
                            world_size=args.world_size, rank=args.rank)
    dist.barrier()

在使用torch.distributed.launch --use_env指令启动时，会自动在python的os.environ中写入RANK、WORLD_SIZE、LOCAL_RANK信息。

在单机多卡的情况下，WORLD_SIZE代表着使用进程数量(一个进程对应一块GPU)，这里RANK和LOCAL_RANK这里的数值是一样的，代表着WORLD_SIZE中的第几个进程（GPU）。
在多机多卡的情况下，WORLD_SIZE代表着所有机器中总进程数(一个进程对应一块GPU)，RANK代表着是在WORLD_SIZE中的哪一个进程，LOCAL_RANK代表着当前机器上的第几个进程（GPU）。

所以在init_distributed_mode函数中会读取os.environ中的参数RANK、WORLD_SIZE、LOCAL_RANK信息。通过读取这些信息，就知道了自己是第几个线程，应该使用哪块GPU设备。通过torch.cuda.set_device()方法设置当前使用的GPU设备。然后使用dist.init_process_group()方法去初始化进程组，其中backend为通信后端，如果使用的是Nvidia的GPU建议使用NCCL；init_method为初始化方法，这里直接使用默认的env://当然也支持TCP或者指像某一共享文件；world_size这里就是该进程组的进