[pytorch]单多机下多GPU下分布式负载均衡训练

Posted wildkid1024

tags:

篇首语:本文由小常识网(cha138.com)小编为大家整理,主要介绍了[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


以上是关于[pytorch]单多机下多GPU下分布式负载均衡训练的主要内容,如果未能解决你的问题,请参考以下文章

jmeter 分布式多机联合负载,主制机不能连接到负载机可能的原因

Pytorch Multi-GPU原理与实现(单机多卡)

[Nutch]Hadoop多机完全分布式模式主机配置

ubantu 下 tomcat + nginx 实现负载均衡

为什么没有 5 层6 层负载均衡?

Pytorch中多GPU训练指北