Pytorch模型在GPU上训练的参数在CPU上移植

希望将训练好的模型加载到新的网络上。如上面题目所描述的，PyTorch在加载之前保存的模型参数的时候，遇到了问题。

Unexpected key(s) in state_dict: "module.features. ...".，Expected ".features....". 直接原因是key值名字不对应。

    表明了加载过程中，期望获得的key值为feature...，而不是module.features....。这是由模型保存过程中导致的，模型应该是在DataParallel模式下面，也就是采用了多GPU训练模型，然后直接保存的。
 

#假设CPU上的模型结构的类为 Class_NetModel
NetModel = Class_NetModel()
print(NetModel.state_dict().keys()) #输出类似odict_keys(['ngram_conv_1_3.0.weight',...

# 根据环境判断是否在GPU上运行
Device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
parallel_model = nn.DataParallel(ClassModel)  # 将模型并行化，以便于使用所有的GPU
ClassModel = parallel_model.to(Device)

modelparam_path = '../model_dump/model.pth'  #这是在GPU上训练出来的参数
checkpoint = torch.load(modelparam_path, map_location=Device)
ClassModel.load_state_dict(checkpoint['model'])
print(ClassModel.state_dict().keys()) #输出类似odict_keys(['module.ngram_conv_1_3.0.weight',...

直接的解决办法就是把GPU上训练的参数的key名称中“module.”删掉。

# original saved file with DataParallel
state_dict = torch.load('checkpoint.pt')  # 模型可以保存为pth文件，也可以为pt文件。
# create new OrderedDict that does not contain `module.`
from collections import OrderedDict
new_state_dict = OrderedDict()
for k, v in state_dict.items():
    name = k[7:] # remove `module.`，表面从第7个key值字符取到最后一个字符，正好去掉了module.
    new_state_dict[name] = v #新字典的key值对应的value为一一对应的值。 
# load params
model.load_state_dict(new_state_dict) # 从新加载这个模型。