Pytorch 加载预训练模型

最新推荐文章于 2024-06-27 12:45:00 发布

徐梓航

最新推荐文章于 2024-06-27 12:45:00 发布

阅读量1.6k

点赞数 2

分类专栏： Pytorch学习文章标签： Pytorch 预训练

本文链接：https://blog.csdn.net/xuyunyunaixuexi/article/details/101906023

版权

本文介绍了如何在PyTorch中加载预训练模型，通过过滤预训练模型字典中不必要的键，覆盖现有模型状态字典的条目，然后加载新的状态字典。此外，还提供了一行代码实现该过程，并引用了多个教程资源。

摘要由CSDN通过智能技术生成

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 ： htt