赞
踩
希望将训练好的模型加载到新的网络上。如上面题目所描述的,PyTorch在加载之前保存的模型参数的时候,遇到了问题。
Unexpected key(s) in state_dict: "module.features. ...".,Expected ".features....".
直接原因是key值名字不对应。
表明了加载过程中,期望获得的key值为feature…,而不是module.features…。这是由模型保存过程中导致的,模型应该是在DataParallel模式下面,也就是采用了多GPU训练模型,然后直接保存的。
You probably saved the model using nn.DataParallel, which stores the model in module, and now you are trying to load it without . You can either add a nn.DataParallel temporarily in your network for loading purposes, or you can load the weights file, create a new ordered dict without the module prefix, and load it back.
解决上面的问题有三个办法:
state_dict = torch.load('checkpoint.pt') # 模型可以保存为pth文件,也可以为pt文件。
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为一一对应的值。
model.load_state_dict(new_state_dict) # 从新加载这个模型。
model.load_state_dict({k.replace('module.',''):v for k,v in torch.load('checkpoint.pt').items()})
#相当于用''代替'module.'。
#直接使得需要的键名等于期望的键名;
如果有多个GPU,将模型并行化,用DataParallel来操作。这个过程会将key值加一个"module. ***"。
model = VGG()# 实例化自己的模型;
checkpoint = torch.load('checkpoint.pt', map_location='cpu') # 加载模型文件,pt, pth 文件都可以;
if torch.cuda.device_count() > 1:
# 如果有多个GPU,将模型并行化,用DataParallel来操作。这个过程会将key值加一个"module. ***"。
model = nn.DataParallel(model)
model.load_state_dict(checkpoint) # 接着就可以将模型参数load进模型。
model = VGGNet()
params=model.state_dict() #获得模型的原始状态以及参数。
for k,v in params.items():
print(k) #只打印key值,不打印具体参数。
features.0.0.weight
features.0.1.weight
features.1.conv.3.weight
features.1.conv.4.num_batches_tracked
model = VGGNet()
checkpoint = torch.load('checkpoint.pt', map_location='cpu')
# Load weights to resume from checkpoint。
# print('**************************************')
# 这个方法能够直接打印出你保存的checkpoint的键和值。
for k,v in checkpoint.items():
print(k)
print("*****************************************")
输出结果为:
module.features.0.0.weight",
"module.features.0.1.weight",
"module.features.0.1.bias
可以看出不匹配,模型的参数中,key值不同,多了module。
————————————————
版权声明:本文为CSDN博主「是否龙磊磊真的一无所有」的原创文章,遵循 CC 4.0 BY-SA 版权协议,转载请附上原文出处链接及本声明。
原文链接:https://blog.csdn.net/qq_32998593/article/details/89343507
Copyright © 2003-2013 www.wpsshop.cn 版权所有,并保留所有权利。