Pytorch DataParallel and DistributedDataParallel

Pytorch DataParallel and DistributedDataParallel

最近试着使用Pytorch跑单机多卡训练,遇到了不少问题,做个总结和教程方便未来观看。我自己也是一个新手,很多东西总结的不好,有问题请多多指教,不懂的地方可以看参考文档,很多东西写的比我详细(本文只针对单机多卡训练,多机多卡训练未经过验证,请酌情观看)
环境:
python 3.7
pytorch 1.4.0

DataParallel

DataParallel是官方最早提供的一个库,使用非常简单,一行就够了

import torch.nn as nn

device = torch.device("cuda")
model = model.to(device)
model = nn.DataParallel(model)

DataParallel的缺点也很明显,只能在单台计算机上使用,而且速度很慢(相较于其他几种加速方式来说),最大的问题在于DataParallel采用了数据分布,但是loss是在某一块卡上计算的,这必然导致会出现比较严重的负载不均衡问题,其他卡用了很少的显存而主卡的显存使用很多,batch size稍微大一点主卡的显存就会炸了,因此更推荐使用其他的方式进行并行计算。

DistributedDataParallel

DistributedDataParallel(DDP)也是官方提供的一个并行计算库,相比于前者,该方式可以更好地进行多机多卡运算,更好的进行并行计算,负载均衡,运行效率也更高,支持的库和方法也更多,虽然使用起来较为麻烦,但对于追求性能来讲是一个更好的选择。

DDP会自动将库分配给n个进程并在n个gpu上运行,loss的计算也是在各个gpu上进行的,很大程度上解决了分配不均衡的问题。

模型设置

  • 调用init设置通信后端
    用于指定通信后端,一般设置nccl就可以,这是nvidia的GPU的参数设定,更具体的设定请参考官方文档。
torch.distributed.init_process_group(backend='nccl')
  • 添加local_rank参数
    local rank是为了DDP运行的时候进行节点进程制定,这个参数设定进去就可以,我们在自己进行计算时不需要人为给定参数,调用启动器的时候会自动设置
parser = argparse.ArgumentParser()
parser.add_argument('--local_rank', default=-

你可能感兴趣的:(Pytorch,python,深度学习,pytorch)