pytorch 多卡负载不均衡

使用3个1080进行模型训练，发现设置batch_size=16把第一张卡给炸掉，但是其他的卡只用了一半不到，严重的负载不均衡啊，因为是租用 极客云GPU这是一个超级扯蛋的平台，我租的4卡1080，开机之后发现才3张卡，还有租双卡的1080Ti，开机之后竟然只有一张卡，😂😂😂😂，现在就想把钱花完，跑路，去其他平台。

这是没进行负载平衡之前的gpu占用，batch_size=16会炸第一张卡，导致整个程序停止。

参考这里的DataParallel imbalanced memory usage方法，把每个GPU的loss分开计算。

• 让每个gpu单独计算loss,再返回即可解决这个问题 即 prediction,loss=model(data,target) (#如果后边不用prediction,连prediction也不用返回，只返回loss)。

• 每个gpu返回一个loss,合到主gpu就是一个list，要loss.mean() 或loss.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
    

def DataParallel_withLoss(model,loss,**kwargs):
    model=FullModel(model, loss)
    if 'device_ids' in kwargs.keys():
        device_ids=kwargs['device_ids']
    else:
        device_ids=None
    if 'output_device' in kwargs.keys():
        output_device=kwargs['output_device']
    else:
        output_device=None
    if 'cuda' in kwargs.keys():
        cudaID=kwargs['cuda'] 
        model=torch.nn.DataParallel(model, device_ids=device_ids, output_device=output_device).cuda(cudaID)
    else:
        model=torch.nn.DataParallel(model, device_ids=device_ids, output_device=output_device).cuda()
    return model

修改之后比未修改好很多了，差一个1G左右，之前是4G。

当然还有其他方法paralle和horovod

使用这两种方法需要修改的代码比较多。

参考：
更大的batch（一）：单GPU、多GPU和分布式系统的实用技巧
pytorch如何使用多块gpu?
增大Batch训练神经网络：单GPU、多GPU及分布式配置的实用技巧
horov
DataParallel imbalanced memory usage
pytorch 多GPU训练总结（DataParallel的使用）
PyTorch使用并行GPU处理数据