# 加载相同state_dict
def load_same_state_dict(myself_model: nn.Module, pretrain_model: nn.Module):
pretrain_dict = pretrain_model.state_dict()
myself_dict = myself_model.state_dict()
# 当模型中的某层是同时在两个模型中共有时才取出
pretrain_dict = {k: v for k, v in pretrain_dict.items() if k in myself_dict}
myself_dict.update(pretrain_dict)
# 然后将参数导入你的模型即可
myself_model.load_state_dict(myself_dict)
return myself_model
重构预训练模型后加载参数
最新推荐文章于 2023-01-15 18:51:29 发布