Pytorch 使用多GPU训练，有用

model

先把model都定义好，再统一放到GPU上，不然可能会报错
不要动optimizer

if len(conf.device_ids)>1:# 多GPU
    self.model = torch.nn.DataParallel(self.model, device_ids=conf.device_ids).cuda()
	self.head = torch.nn.DataParallel(self.head, device_ids=conf.device_ids).cuda()
	#self.optimizer = torch.nn.DataParallel(self.optimizer, device_ids=conf.device_ids).cuda()
else: # 单GPU
	self.model = self.model.to(conf.device)
	self.head = self.head.to(conf.device)
	#self.optimizer = self.optimizer.to(conf.device)
	      

data

imgs = imgs.cuda()
labels = labels.cuda()

一些BUG

有些帖子说在取用state_dict和param_groups之前要加入.module，我没有改动这一块。

---

参考博客
参考代码