pytorch的多gpu训练

1.多GPU训练，出现out of memory

出现情景：预训练模型是使用gpu0训练得到，然后要在多gpu的服务器上进行微调，使用gpu id为[4,5,6,7]，然后出现报错如下：

cuda runtime error (2) : out of memory at /pytorch/aten/src/THC/THCTensorRandom.cu:25

错误原因分析：在加载预训练模型的位置报错，在台式机（只有一块gpu）将模型加载打印输出参数：

代码：

checkpoint = torch.load("/home/final.pth")
for k, v in checkpoint.items():
     print(k)
     print(v)

打印输出出现：

 

发现模型加载的时候就将参数传入到gpu当中，而我在服务器上出现报错应该是模型参数直接加载到了gpu id = 0 的显卡上，但是gpu id= 0的显卡上显存已经满了；所以报出了 内存溢出的问题；但是我在操作的时候，直接使用gpu id = 4的显卡也会出现错误，出现错误的模型加载代码如下：

if conf.pretrained == True:
      checkpoint = torch.load("./model_mobilefacenet.pth",
map_location={'cuda:1':'cuda:0'})
      self.model_mobile.load_state_dict(checkpoint)

然后修改成这样也出错：

if conf.pretrained == True:
      checkpoint = torch.load("./model_mobilefacenet.pth",
,map_location=lambda storage, lo