PyTorch用预训练模型初始化网络的部分权重

PyTorch用预训练模型初始化网络的部分权重

关于从预训练的模型加载参数时可能会遇到的一些阻碍,笔者对自己的绵薄经验在此进行分享,也便于自己之后的温习。

加载初始化参数可参照论坛中 apaszke的做法,其删除了与当前model不匹配的key。
代码为:

pretrained_dict = ...
model_dict = model.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)
# 3. load the new state dict
model.load_state_dict(model_dict)

如果修改后的模型与预训练模型中无key重名现象,此法可以使用。但是如果有重名现象,则可能出现训练出错,如加载参数与模型的正确维度不匹配等。此时,加一个尺寸相等的限制条件即可。代码如下:

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

pretrained_dict = {k: v for k, v in pretrained_dict.items() if k in model_dict 
   					and v.size() == model_dict[k].size()}
model_dict.update(pretrained_dict)
model.load_state_dict(model_dict)

此方法参考了论坛中 ImgPrcSng的做法。

  • 1
    点赞
  • 15
    收藏
    觉得还不错? 一键收藏
  • 0
    评论
评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值