Pytorch 加载部分模型参数样例

加载相同名字的参数

pretrained_dict=torch.load(model_weight)
model_dict=myNet.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)
myNet.load_state_dict(model_dict)

其他加载方式可修改中间的列表生成式来过滤部分层

例如过滤FC层

para_state_dict = torch.load("./cmt_small.pth")["model"]
model_dict = model.state_dict()
pretrained_dict = {k: v for k, v in para_state_dict.items() if k in model_dict and "head" not in k}
model_dict.update(pretrained_dict)
model.load_state_dict(model_dict)