可能会遇到有一个多gpu的训练后保存的模型,但是后续要在单gpu的机子上使用,在torxh.load时报错。
这是因为nn.DataParallel会在模型参数结构前面加一个module.
比如你是这么save的
state = {'epoch': epoch, 'state_dict': self.model.state_dict(),
'optimizer': self.optimizer.state_dict(),
'info': self.info,
'scheduler': self.scheduler.state_dict()
}
w_dict = torch.load(path + "/pretrain_model_name", map_location=lambda storage, loc: storage) # 权重是 多gpu的,去掉module字样
# print( "original dict:--", w_dict)
new_state_dict= OrderedDict()
for k,v in w_dict['state_dict'].items(): # 字典 'state_dict':
namekey= k[7:] if k.startswith('module.') else k
new_state_dict[namekey]= v
print("new_state_dict---: {}".format(new_state_dict) )
model.load_state_dict(new_state_dict)