PyTorch加载模型model.load_state_dict()问题,Unexpected key(s) in state_dict: “module.features..,Expected .
希望將訓練好的模型加載到新的網絡上。如上面題目所描述的,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.
解決上面的問題有三個辦法:
1. 對load的模型創建新的字典,去掉不需要的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) # 從新加載這個模型。
2. 直接用空白''代替'module.'
model.load_state_dict({k.replace('module.',''):v for k,v in torch.load('checkpoint.pt').items()})# 相當于用''代替'module.'。
#直接使得需要的鍵名等于期望的鍵名。
3. 最簡單的方法,加載模型之后,接著將模型DataParallel,此時就可以load_state_dict。
如果有多個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進模型。
?4. 總結
? ? 從出錯顯示的問題就可以看出,key值不匹配,因此可以選擇多種方法,將模型參數加載進去。 這個方法通常會在load_state_dict過程中遇到。將訓練好的一個網絡參數,移植到另外一個網絡上面,繼續訓練?;蛘邔⒂柧毢玫木W絡checkpoint加載進模型,再次進行訓練。可以打印出model state_dict來看出兩者的差別。
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。
PS: 2020-12-25
在移植參數的過程中,對于出現?.total_ops和.total_params結尾的參數,可參考以下代碼:
from collections import OrderedDict
checkpoint = torch.load(pretrained_model_file_path,map_location=(None if use_cuda and not remap_to_cpu else "cpu"))
new_state_dict = OrderedDict()
for k, v in checkpoint.items():if not k.endswith('total_ops') and not k.endswith('total_params'):name = k[7:]new_state_dict[name] = v
如果有用,記得點贊👍加收藏哦。!!!!
總結
以上是生活随笔為你收集整理的PyTorch加载模型model.load_state_dict()问题,Unexpected key(s) in state_dict: “module.features..,Expected .的全部內容,希望文章能夠幫你解決所遇到的問題。
- 上一篇: 多囊卵巢怀二胎好怀吗
- 下一篇: 基于FCN,U-Net的深度学习医学影像