pytorch 分布式训练初探

2024年5月28日 291次阅读 来源: lbin

pytorch 分布式训练初探

为什么需要分布式

众所周知，深度神经网络发展到现阶段，离不开GPU和数据。经过这么多年的积累，GPU的计算能力越来越强，数据也积累的越来越多，大家会发现在现有的单机单卡或者单机多卡上很难高效地复现模型，甚至对于有些新的数据集来讲，单机训练简直就是噩梦。

DatasetImagesMS COCO115,000Open Image dataset v41,740,000

为什么单机8卡也会是噩梦呢？我们拿COCO和Google最近Release出来的Open Image dataset v4来做比较，训练一个resnet152的检测模型，在COCO上大概需要40个小时，而在OIDV4上大概需要40天，这还是在各种超参数正确的情况下，如果加上调试的时间，可能一个模型调完就该过年了吧。

所以这个时候我们需要分布式。

pytorch 分布式简介

PyTorch分布式功能在0.4版本中得到大幅完善，它允许在多台机器之间交换Tensors，这样就可以通过多台机器和更大的minibatch扩展网络训练。

torch.distributed provides an MPI-like interface for exchanging tensor data across multi-machine networks. It supports a few different backends and initialization methods.

Currently torch.distributed supports four backends, each with different capabilities. The table below shows which functions are available for use with CPU / CUDA tensors. MPI supports cuda only if the implementation used to build PyTorch supports it.

pytorch文档里有这么一段介绍，但其实可选的只有Gloo。它是一个类似MPI的通信库，你不需要考虑内存数据的拷贝，只需要实现逻辑就可以。

这里引入了一个新的函数model = torch.nn.parallel.DistributedDataParallel(model)为的就是支持分布式模式

不同于原来在multiprocessing中的model = torch.nn.DataParallel(model,device_ids=[0,1,2,3]).cuda()函数，这个函数只是实现了在单机上的多GPU训练，根据官方文档的说法，甚至在单机多卡的模式下，新函数表现也会优于这个旧函数。

这里要提到两个问题：

  • 每个进程都有自己的Optimizer同时每个迭代中都进行完整的优化步骤，虽然这可能看起来是多余的，但由于梯度已经聚集在一起并跨进程平均，因此对于每个进程都是相同的，这意味着不需要参数广播步骤，从而减少了在节点之间传输张量tensor所花费的时间。
  • 另外一个问题是Python解释器的，每个进程都包含一个独立的Python解释器，消除了来自单个Python进程中的多个执行线程，模型副本或GPU的额外解释器开销和“GIL-thrashing”。 这对于大量使用Python运行时的模型尤其重要。

初始化

torch.distributed.init_process_group(backend, init_method='env://', **kwargs)

参数说明：

  • backend(str): 后端选择，包括上面那几种 tcp mpi gloo
  • init_method(str，optional): 用来初始化包的URL, 我理解是一个用来做并发控制的共享方式
  • world_size(int, optional): 参与这个工作的进程数
  • rank(int,optional): 当前进程的rank
  • group_name(str,optional): 用来标记这组进程名的

解释一下init_method()也有这三种方式，具体可参看http://pytorch.org/docs/master/distributed.html

  • file:// 共享文件系统（要求所有进程可以访问单个文件系统）有共享文件系统可以选择
  • tcp:// IP组播（要求所有进程都在同一个网络中）比较好理解，不过需要手动设置rank
  • env:// 环境变量（需要您手动分配等级并知道所有进程可访问节点的地址）默认是这个

分布式 Hello World

https://github.com/pytorch/examples/tree/master/imagenet
这里，常规的操作就不多叙述了，主要讲一下和分布式相关的代码部分。

parser.add_argument('--world-size', default=2, type=int, help='number of distributed processes')
parser.add_argument('--dist-url', default='tcp://172.16.1.186:2222', type=str, help='url used to set up distributed training')
parser.add_argument('--dist-backend', default='gloo', type=str, help='distributed backend')
parser.add_argument('--dist-rank', default=0, type=int, help='rank of distributed processes')

这几个是必要的参数设置，其中最后一个是官网没有的

if args.distributed:
    dist.init_process_group(backend=args.dist_backend, init_method=args.dist_url,world_size=args.world_size,rank=args.dist_rank)

这个是分布式的初始化，同样，最后添加一个rank

model.cuda()
model = torch.nn.parallel.DistributedDataParallel(model)

这里，把我们平时使用的单机多卡，数据并行的API

model = torch.nn.DataParallel(model).cuda()

换掉即可。

if args.distributed:
        train_sampler = torch.utils.data.distributed.DistributedSampler(train_dataset)

最后使用这个官方给的划分方法，把数据集划分即可。

需要注意的点

  • 通信主机地址一定要写ib的ip，否则带宽不够会造成Gloo底层触发延迟错误
  • 一定要加rank参数
  • train_dataset最好不要用自己写的sampler，否则还需要再实现一遍分布式的数据划分方式

后续的计划

一次讲完分布式太难，所以这会是一个系列，系列结束后希望可以让你自如地训练自己的检测、分割或者各种各样的模型。

系列的内容包括但不限于:

  • 使用uber/horovod加速训练
  • 详解Train ImageNet In one Hour
  • 如何得到更好的ResNet50
  • 训练检测网络 or Mask RCNN

敬请期待~

    原文作者：lbin
    原文地址: https://zhuanlan.zhihu.com/p/43424629
    本文转自网络文章，转载此文章仅为分享知识，如有侵权，请联系博主进行删除。
点赞