Pytorch单GPU、多GPU训练的几个细节对比

Pytorch单GPU、多GPU训练的区别主要在三个地方：训练前指定GPU、训练过程中保存模型和加载刚刚保存的模型。

训练前指定GPU

单GPU：

os.environ["CUDA_VISIBLE_DEVICES"] = '0'

多GPU：

os.environ["CUDA_VISIBLE_DEVICES"] = '0,1'
if torch.cuda.device_count() > 1:
    print("Let's use", torch.cuda.device_count(), "GPUs!")
    self.model = nn.DataParallel(self.model)

训练过程中保存模型

单GPU：

state = {'model': self.model.state_dict(), 'epoch': ite}
torch.save(state, self.model.name())

多GPU：

if isinstance(self.model,torch.nn.DataParallel):##判断是否并行
    self.model = self.model.module
    
state = {'model': self.model.state_dict(), 'epoch': ite}
torch.save(state, self.model.name())

if torch.cuda.device_count() > 1:
    self.model = nn.DataParallel(self.model)

多加的上面两句是为了解决下面的问题

AttributeError: 'DataParallel' object has no attribute 'name'

如果不加最后两句，也不会报错，但是后面训练都会变成单GPU，也就是会导致下面的结果。（我用的两个GPU）

加上后两句之后：

需要注意前两句、后两句以及原来两句的相对位置不能颠倒，例如把原来的第一句放到最前面，在后面加载模型的时候可能会出现问题。

加载刚刚保存的模型

单GPU：

checkpoint = torch.load(self.model.name())
self.model.load_state_dict(checkpoint['model'])

多GPU改成：

if isinstance(self.model,torch.nn.DataParallel):    
    self.model = self.model.module
if torch.cuda.is_available(): #gpu
    checkpoint = torch.load(self.model.name())
else: #cpu
    checkpoint = torch.load(self.model.name(),map_location=lambda storage, loc: storage)
self.model.load_state_dict(checkpoint['model'])