Pytorch 加载预训练模型

2022-10-25 11:28:15

pretrained_dict =...

model_dict = model.state_dict()

# 1. filter out unnecessary keys
pretrained_dict = {k: v for k, vin pretrained_dict.items() if k inmodel_dict}


# 2. overwrite entries in the existing state dict

model_dict.update(pretrained_dict)


# 3. load the new state dict

model.load_state_dict(model_dict)
例如:

myNet=resnet18()
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)

也可以用一句话搞定:myNet.load_state_dict(torch.load(model_weight),strict=False)

ref :https://blog.csdn.net/LXX516/article/details/80124768

https://blog.csdn.net/lscelory/article/details/81482586

https://blog.csdn.net/strive_for_future/article/details/83240081

https://www.cnblogs.com/wmlj/p/9917827.html

  • 作者:徐梓航
  • 原文链接:https://blog.csdn.net/xuyunyunaixuexi/article/details/101906023
    更新时间:2022-10-25 11:28:15