模型B加载模型A部分结构(B)的权重参数,保证机构名与权重名对齐加载方法

 
 # (本例子情况:A中的(B)的权重key比B的结构名中多了'res.' 可类比修改)

# 模型A中存在的结构B权重文件路径:
weight = '权重路径'       
# 模型B:以res50为例
model = torchvision.models.resnet50(num_classes).to(device)

# 建一个空字典,模拟权重文化
state_dict = dict()

# 把模型A的权重中  key-value 的 key 打印出来
print('[pretrained_model_keys]:', *torch.load(weight).keys(), sep="\n")

for k, v in torch.load(weight).items():
# 将多余出来的res. 替换为空
    state_dict[k.replace('res.', '')] = v
# 加载即可
# model.load_state_dict(torch.load(weight), strict=False)

# missing_keys代表该模型中没有匹配到权重的层
# unexpected_keys代表 所加载的权重文件中,未匹配到层结构的权重名
missing_keys, unexpected_keys = model.load_state_dict(state_dict, strict=False)
print('[missing_keys]:', *missing_keys, sep="\n")
print('[unexpected_keys]:', *unexpected_keys, sep="\n")

                
  • 1
    点赞
  • 0
    收藏
    觉得还不错? 一键收藏
  • 1
    评论

“相关推荐”对你有帮助么?

  • 非常没帮助
  • 没帮助
  • 一般
  • 有帮助
  • 非常有帮助
提交
评论 1
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

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

抵扣说明:

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

余额充值