如何读取部分的预训练模型

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