pytorch 分布式多卡训练DistributedDataParallel 踩坑记

前言:

  近几天想研究了多卡训练,就花了点时间,本以为会很轻松,可是好多坑,一步一步踏过来,一般分布式训练分为单机多卡与多机多卡两种类型;

   主要有两种方式实现:

1、DataParallel:   Parameter Server模式,一张卡位reducer,实现也超级简单,一行代码

   DataParallel是基于Parameter server的算法,负载不均衡的问题比较严重,有时在模型较大的时候(比如bert-large),reducer的那张卡会多出3-4g的显存占用,也就是说用到DataParallel的时候gpu 0的使用内存会大大超出其他显卡的使用内存。

2、DistributedDataParallel:官方建议用新的DDP,采用all-reduce算法,本来设计主要是为了多机多卡使用,但是单机上也能用

  为什么要分布式训练?

  1. 可以用多张卡,总体跑得更快
  2. 可以得到更大的 BatchSize
  3. 有些分布式会取得更好的效果

分布式原理性解释:

https://www.cnblogs.com/yh-blog/p/12877922.html

https://zhuanlan.zhihu.com/p/113694038

主要分为以下几个部分:

  1. 单机多卡,DataParallel(最常用,最简单)
  2. 单机多卡,DistributedDataParallel(较高级)、多机多卡,DistributedDataParallel(最高级)
  3. 如何启动训练
  4. 模型保存与读取
  5. 注意事项

一、单机多卡(DataParallel)

from torch.nn import DataParallel

device = torch.device("cuda")
#或者device = torch.device("cuda:0" if True else "cpu")

model = MyModel()
model = model.to(device)
model = DataParallel(model)
#或者model = nn.DataParallel(model,device_ids=[0,1,2,3])

比较简单,只需要加一行代码就行, model = DataParallel(model)

二、多机多卡、单机多卡(DistributedDataParallel)

建议先把注意事项看完在修改代码,防止出现莫名的bug,修改训练代码如下:

其中opt.local_rank要在代码前面解析这个参数,可以去后面看我写的注意事项;

    from torch.utils.data.distributed import DistributedSampler
    import torch.distributed as dist
    import torch

    # Initialize Process Group
    dist_backend = 'nccl'
    print('args.local_rank: ', opt.local_rank)
    torch.cuda.set_device(opt.local_rank)
    dist.init_process_group(backend=dist_backend)

    model = yourModel()#自己的模型
    device = torch.device(opt.local_rank)
    model.to(device)#经后来测试,好像注释这一条代码也可以正常训练;
    if torch.cuda.device_count() > 1:
        print("Let's use", torch.cuda.device_count(), "GPUs!")
        # 5) 封装
        # model = torch.nn.parallel.DistributedDataParallel(model,
        #                                                   device_ids=[opt.local_rank],
        #                                                   output_device=opt.local_rank)
        model = torch.nn.parallel.DistributedDataParallel(model.cuda(), device_ids=[opt.local_rank])

    dataset = ListDataset(train_path, augment=True, multiscale=opt.multiscale_training, img_size=opt.img_size, normalized_labels=True)#自己的读取数据的代码
    world_size = torch.cuda.device_count()
    datasampler = DistributedSampler(dataset, num_replicas=dist.get_world_size(), rank=opt.local_rank)

    dataloader = torch.utils.data.DataLoader(
        dataset,
        batch_size=opt.batch_size,
        shuffle=False,
        num_workers=opt.n_cpu,
        pin_memory=True,
        collate_fn=dataset.collate_fn,
        sampler=datasampler
    )#在原始读取数据中加sampler参数就行


.....

训练过程中,数据转cuda
      imgs = imgs.to(device)
      targets = targets.to(device)

三、如何启动训练

  1、DataParallel方式

正常训练即可,即

python3  train.py

       2、DistributedDataParallel方式

          需要通过torch.distributed.launch来启动,一般是单节点,

CUDA_VISIBLE_DEVICES=0,1 python3 -m torch.distributed.launch --nproc_per_node=2 train.py

          其中CUDA_VISIBLE_DEVICES 设置用的显卡编号,--nproc_pre_node 每个节点的显卡数量,一般有几个显卡就用几个显卡

当然还可以多节点训练,这个没了解清楚,想了解的可以戳这个

https://zhuanlan.zhihu.com/p/113694038

多节点

python3 -m torch.distributed.launch --nproc_per_node=NUM_GPUS_YOU_HAVE --nnodes=2 --node_rank=0
#两个节点,在0号节点

要是训练成功,就会打印出几个信息,有几个卡就打印几个信息,如下图所示:

pytorch 分布式多卡训练DistributedDataParallel 踩坑记_第1张图片

四、模型保存与读取

以下a、b是对应的,用a保存,就用a方法加载

  1、保存

    a、只保存参数

torch.save(model.module.state_dict(), path)

    b、保存参数与网络

torch.save(model.module,path)

  2、加载

          a、多卡加载模型预训练;

model = Yourmodel()
if opt.pretrained_weights:
        if opt.pretrained_weights.endswith(".pth"):
            model.load_state_dict(torch.load(opt.pretrained_weights))
        else:
            model.load_darknet_weights(opt.pretrained_weights)

         单卡加载模型,需要加载模型时指定主卡读模型,而且这个'cuda:0',是看你训练的模型是0还是1(否则就会出错RuntimeError: Attempting to deserialize object on CUDA device 1 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device),可以根据自己的更改:

model = Yourmodel()
if opt.pretrained_weights:
        if opt.pretrained_weights.endswith(".pth"):
            model.load_state_dict(torch.load(opt.pretrained_weights,map_location="cuda:0"))
        else:
            model.load_darknet_weights(opt.pretrained_weights)

b、单卡加载模型;

  同样也要指定读取模型的卡。  

model = torch.load(opt.weights_path, map_location="cuda:0")

多卡加载预训练模型,以b这种方式还没跑通。 

模型加载其他参考:https://www.zhihu.com/question/67726969

五、注意事项

   1、model后面添加module

    获取到网络模型后,使用并行方法,并将网络模型和参数移到GPU上。注意,若需要修改网络模块或者获得模型的某个参数,一定要在model后面加上.module,否则会报错,比如:

   model.img_size  要改成  model.module.img_size

   2、.cuda或者.to(device)等问题

              device是自己设置,如果.cuda出错,就要化成相应的device

    model(如:model.to(device))
    input(通常需要使用Variable包装,如:input = Variable(input).to(device))
    target(通常需要使用Variable包装
    nn.CrossEntropyLoss()(如:criterion = nn.CrossEntropyLoss().to(device))
   3、args.local_rank的参数

    通过torch.distributed.launch来启动训练,torch.distributed.launch 会给模型分配一个args.local_rank的参数,所以在训练代码中要解析这个参数,也可以通过torch.distributed.get_rank()获取进程id。

parser.add_argument("--local_rank", type=int, default=-1, help="number of cpu threads to use during batch generation")

    4、报错:RuntimeError: Attempting to deserialize object on CUDA device 2 but torch.cuda.device_count() is 1

  • 原因:在使用Pytorch加载模型时报错。加载的模型是用两个GPU训练的,而加载模型的电脑只有一个GPU,所以会出错。
  • 解决:model = torch.load(model_path)
    改为:model = torch.load(model_path, map_location='cuda:0')
    如果是4块到2块:就把map_location改为:map_location={'cuda:1': 'cuda:0'}

 

你可能感兴趣的:(Pytorch)