[pytorch] 模型加载，torch.load , keys not match

可能会遇到有一个多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)