希望将训练好的模型加载到新的网络上。如上面题目所描述的,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) # 从新加载这个模型。