pytorch加载预训练模型中的部分参数

有时候我们构建的模型跟预训练模型的网络结构是不相同，为了将预训练模型中与我们模型重合的部分加载进来，就需要用到加载部分参数的方法，比如说

预训练模型结构为：ResNet18 -> fc 

我们要构建的模型结构为 ResNet -> fc1 -> fc2

我们希望把在预训练模型中训练好的ResNet18加载进新模型的ResNet18里（因为不想加载预训练模型的fc层，所以新模型里最好不要重复命名fc层，但其实即使重复命名，只要不要结构也相同就行了）：

model=NewModel()#构建的新模型
model_dict = model.state_dict()
for k, v in model_dict.items(): 
    print(k)

pretrained_dict = torch.load('xxx.pkl')
for k, v in pretrained_dict.items():
    print(k)
pretrained_dict = {k: v for k, v in pretrained_dict.items() if k in model_dict}
#取出预训练模型中与新模型的dict中重合的部分


model_dict.update(pretrained_dict)#用预训练模型参数更新new_model中的部分参数
model.load_state_dict(model_dict) #将更新后的model_dict加载进new model中

更详细的pytorch模型保存和加载问题，参考我的另一个博客https://blog.csdn.net/weixin_41519463/article/details/103205665