Pytorch单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'])