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