pytorch使用DataParallel进行GPU并行训练，以及AttributeError: 'DataParallel' object has no attribute '**'错误的解决

目录

请让模型奔跑

一个bug出现了

---

请让模型奔跑

首先由于GPU暂时没有爆显存的担心，主要是想把程序尽快跑起来，所以参考了https://zhuanlan.zhihu.com/p/86441879的方法，在他的文章中，不仅对单机多GPU的情况进行了介绍，同时也给出了平衡显存的方法，以及多机多GPU的方法，并且介绍很简洁易懂。

在我的程序中，主要用到的是最简单的情况，步骤如下：

• 使用单机多卡进行训练，设置程序可见的GPU编号：

  os.environ['CUDA_VISIBLE_DEVICES'] = '0,3'
  

   

• 对模型执行以下代码：

  model = nn.DataParallel(model)  #我们自定义的模型
  model = model.cuda()

   

• 对数据执行以下代码：

  inputs = inputs.cuda()   #我们模型的输入数据
  labels = labels.cuda()   #我们数据的标签

   

• 上面的博文作者也提到，以上代码与pytorch官方代码具有相同的作用：

  #device为用户主机上从0号开始可用的GPU编号
  device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
  
  
  #Model为用户的网络模型
  model = Model(input_size, output_size)
  
  if torch.cuda.device_count() > 1:
    print("Let's use", torch.cuda.device_count(), "GPUs!")
    # dim = 0 [30, xxx] -> [10, ...], [10, ...], [10, ...] on 3 GPUs
    model = nn.DataParallel(model)
  
  model.to(device)
  

   

一个bug出现了

在跑程序的时候出现了一个bug提示：AttributeError: 'DataParallel' object has no attribute '**'

这次参考了：https://github.com/jytime/Mask_RCNN_Pytorch/issues/2

其中提到，因为使用torch.nn.DataParallel而造成了AttributeError: 'DataParallel' object has no attribute '**'，‘**’为用户自定义的模型中的函数，此时因为使用了DataParallel函数对原本用户定义的model进行了包装，原model变为了DataParallel的一个子模块。

因此想要再次调用原model中的函数，需要将原本的model.load()函数变为model.module.load()形式。其中load()为我在model中自定义的函数，综上，在model与函数中间加入.module即可。