加载预训练权重的部分参数

最新推荐文章于 2024-07-26 21:47:55 发布

梦坠凡尘(AICV与前沿)

最新推荐文章于 2024-07-26 21:47:55 发布

阅读量1.2k

点赞数 1

分类专栏： Pytorch 深度学习文章标签：深度学习 Pytorch

本文链接：https://blog.csdn.net/c2250645962/article/details/108208526

版权

深度学习同时被 2 个专栏收录

74 篇文章 14 订阅

订阅专栏

Pytorch

16 篇文章 33 订阅

订阅专栏

pretrained_dict = ...
model_dict = model.state_dict()

# pretrained_dict: ['A', 'B', 'C', 'D']
# model_dict: ['A', 'B', 'C', 'E']


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

# pretrained_dict: ['A', 'B', 'C']
# model_dict: ['A', 'B', 'C', 'E']

# 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)