載入模型

轉載

1. 讀取預訓練模型和現有模型的重合部分

reference: https://discuss.pytorch.org/t/how-to-load-part-of-pre-trained-model/1113/3

pretrained_dict = ...
model_dict = model.state_dict()

# 1. filter out unnecessary keys
pretrained_dict = {k: v for k, v in pretrained_dict.items() if k in model_dict}
# 2. overwrite entries in the existing state dict
model_dict.update(pretrained_dict) 
# 3. load the new state dict
model.load_state_dict(pretrained_dict)

 

2. 如果預訓練模型有Module而目前的沒有

參考：https://blog.csdn.net/kaixinjiuxing666/article/details/85115077

# original saved file with DataParallel
state_dict = torch.load('myfile.pth')
# 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.`
    new_state_dict[name] = v
# load params
model.load_state_dict(new_state_dict)

或者先宣告parrallel再讀取預訓練模型

 

3. 如果目前的模型只有部分需要讀取預訓練模型，別的部分（比如自己新增的）在預訓練模型中並沒有，則應該使用strict=False標誌，即：

def load_state_dict(self, state_dict, strict=True):

參考：https://pytorch.org/docs/master/_modules/torch/nn/modules/module.html#Module.load_state_dict