zoukankan      html  css  js  c++  java
  • [pytorch]单多机下多GPU下分布式负载均衡训练

    说明

    在前面讲模型加载和保存的时候,在多GPU情况下,实际上是挖了坑的,比如在多GPU加载时,GPU的利用率是不均衡的,而当时没详细探讨这个问题,今天来详细地讨论一下。

    问题

    在训练的时候,如果GPU资源有限,而数据量和模型大小较大,那么在单GPU上运行就会极其慢的训练速度,此时就要使用多GPU进行模型训练了,在pytorch上实现多GPU训练实际上十分简单:

    只需要将模型使用nn.DataParallel进行装饰即可。

    model = nn.DataParallel(model,device_ids=range(torch.cuda.device_count()))
    

    但是问题在于这样直接处理后的模型的负载可能是不均衡的,因为在不同的GPU上进行运算,而最后的loss计算过程是要合并到主GPU上,这样主GPU的的占用率将比较高,而其余GPU的利用率则没有那么高。

    解决思路

    1. 一种比较简单的解决方法是将模型计算loss的过程封装在model里,这样每个GPU forward的时候就会计算到对应的loss,当然缺陷在于这样每次得到的loss都是一个数组,需要另外mean或者sum处理一下。
    class FullModel(nn.Module):
      def __init__(self, model, loss):
        super(FullModel, self).__init__()
        self.model = model
        self.loss = loss
    
      def forward(self, targets, *inputs):
        outputs = self.model(*inputs)
        loss = self.loss(outputs, targets)
        return torch.unsqueeze(loss,0),outputs
    

    在上述的代码中,构建了另外一个包含model和loss的壳,在壳里计算loss的值,需要注意的是,在进行DataParallel时,也需要对这个并行,而到了收集loss的时候,则使用loss的和:

    loss,_ = model(gt,input)
    loss = loss.sum()
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()
    

    已经有人造了这个轮子,并开源了出来,可以参考:https://github.com/zhanghang1989/PyTorch-Encoding 代码库,整个写法依然没有太大的变化:

    from utils.encoding import DataParallelModel, DataParallelCriterion
    model = DataParallelModel(model)
    criterion = DataParallelCriterion(criterion)
    
    1. 通过distributedDataparallel来实现

    实际上官方考虑过负载不均衡的问题,在文档中也推荐使用distributedDataparallel(ddp)进行训练,尽管ddp是用来解决不同机器的分布式训练问题的。

    This is the highly recommended way to useDistributedDataParallel, with multiple processes, each of which operates on a single GPU. This is currently the fastest approach to do data parallel training using PyTorch and applies to both single-node(multi-GPU) and multi-node data parallel training. It is proven to be significantly faster thantorch.nn.DataParallelfor single-node multi-GPU data parallel training.

    ddp使用起来比DataParallel更快,数据也更均衡,但是缺点是配置起来相对要麻烦一些。

    # 初始化使用的后端
    torch.distributed.init_process_group(backend="nccl")
    
    # 对数据进行划分
    train_sampler = torch.utils.data.distributed.DistributedSampler(train_dataset)
    test_sampler = torch.utils.data.distributed.DistributedSampler(test_dataset)
    train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=batch_size, shuffle=False, num_workers=n_worker, pin_memory=True, sampler=train_sampler)
    test_loader = torch.utils.data.DataLoader(test_dataset, batch_size=batch_size, shuffle=False, num_workers=n_worker, pin_memory=True, sampler=test_sampler) # sampler和shuffle不能同时使用
    
    model=torch.nn.parallel.DistributedDataParallel(model)
    

    注意:需要注意的是,尽量设定pin_memory参数为true,该参数是锁存操作,使用会加快数据读取速度,但是此时要限定内存的大小是要使用显存的两倍
    以上就配置好了,经过测试,使用ddp的训练时间比DataParallel快一倍。
    在运行的时候,使用以下命令进行分布式训练: python -m torch.distributed.launch --nproc_per_node=NUM_GPUS_YOU_HAVE yourscript.py

    引用

    1. https://discuss.pytorch.org/t/dataparallel-imbalanced-memory-usage/22551/20
    2. https://zhuanlan.zhihu.com/p/95700549
  • 相关阅读:
    「BZOJ2721」「LuoguP1445」 [Violet]樱花(数论
    「USACO08DEC」「LuoguP2921」在农场万圣节Trick or Treat on the Farm(tarjan
    「HNOI2008」「LuoguP3197」越狱(数论
    「CF779B」「LOJ#10201.」「一本通 6.2 练习 4」Sherlock and His Girlfriend(埃氏筛
    「LOJ#10072」「一本通 3.2 例 1」Sightseeing Trip(无向图最小环问题)(Floyd
    「LOJ#10068」「一本通 3.1 练习 3」秘密的牛奶运输(次小生成树
    「USACO15FEB」「LuoguP3121」审查(黄金)Censoring (Gold)(AC自动机
    「LOJ#10056」「一本通 2.3 练习 5」The XOR-longest Path (Trie
    「LOJ#10051」「一本通 2.3 例 3」Nikitosh 和异或(Trie
    「UVA644」 Immediate Decodability(Trie
  • 原文地址:https://www.cnblogs.com/wildkid1024/p/13155377.html
Copyright © 2011-2022 走看看