model.load_state_dict(state_dict)构造模型和预训练模型的结构不一样的情况

预训练的模型可以很好的学到特征,我们想利用预训练的模型再加上分类器,实现分类任务,需要在自己的模型上添加新的层,直接使用 model.load_state_dict(state_dict) 会报错:

RuntimeError: Error(s) in loading state_dict for Model:
    Missing key(s) in state_dict

可以逐个的查找自己的模型如果和预训练的模型有相同的参数名称,那么就复制过来,这样就可以避免直接全部加载出现的错误。

如果预训练模型中有多余的网络结构,可以用del 删去。

# 自己的模型
model = Model()

#加载预训练的模型
checkpoint = torch.load('checkpoint_{:04d}.pth.tar'.format(args.epochs))
state_dict = checkpoint['state_dict']

#删去不必要的结构
for k in list(state_dict.keys()):
    # 删去预训练模型中mlp开头的所有结构
    if k.startswith('mlp'):
          del state_dict[k]

for name, param in model.named_parameters():
    # print(name,param.size())
    # 将所有不是predictor开头的结构冻结
    if not (name.startswith('predictor')):
        param.requires_grad = False

    # 查看model初始化的参数
    print("original param: {} {} ".format(name,param))
    # 从预训练模型中加载    
    if name in list(state_dict.keys()):
        param = state_dict[name]
        print("new param from checkpoint: {} {} ".format(name,param))



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

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

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

抵扣说明:

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

余额充值