加载部分预训练参数方法【解决报错】IndexError:list index out of range

【学习参考】Pytorch 加载、查看预训练模型参数、使用部分预训练模型参数初始化网络

  • 取出自己网络的参数字典   
  • 加载预训练网络的参数字典                               
  • 取出预训练网络的参数字典
  • 自己网络和预训练网络结构一致的层,使用预训练网络对应层的参数初始化
# 取出自己网络的参数字典
model_dict = model.state_dict()     

# 加载预训练网络的参数字典                               
pretrained_dict = torch.load("xxxxxx.pth")

# 取出预训练网络的参数字典
keys = []
for k, v in pretrained_dict.items():
       keys.append(k)
i = 0
 
# 自己网络和预训练网络结构一致的层,使用预训练网络对应层的参数初始化
for k, v in model_dict.items():
    if v.size() == pretrained_dict[keys[i]].size():
         model_dict[k] = pretrained_dict[keys[i]]
         #print(model_dict[k])
         i = i + 1
model.load_state_dict(model_dict)

2【报错】IndexError:list index out of range

检查list len()发现新创建的模型参数有898;预训练的模型参数有368 。

 思路:反过来遍历!

# 取出自己网络的参数字典
model_dict = model.state_dict()     

# 加载预训练网络的参数字典                               
pretrained_dict = torch.load("xxxxxx.pth")

# 取出预训练网络的参数字典
keys = []
for k, v in pretrained_dict.items():
       keys.append(k)
i = 0
 
# 自己网络和预训练网络结构一致的层,使用预训练网络对应层的参数初始化
for k, v in model_dict.items():
    if v.size() == model_dict[keys[i]].size():
         model_dict[keys[i]]=pretrained_dict[k] 
         #print(model_dict[k])
         i = i + 1
model.load_state_dict(model_dict)

3 其他的调用的方式

if path is not None:
     self.load(path)
import torch

class BaseModel(torch.nn.Module):
    # path为pt文件路径:
    def load(self, path):
        """Load model from file.

        Args:
            path (str): file path
        """
        parameters = torch.load(path, map_location=torch.device("cpu"))

        if "optimizer" in parameters:
            parameters = parameters["model"]

        self.load_state_dict(parameters)

评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包

打赏作者

MengYa_DreamZ

你的鼓励将是我创作的最大动力

¥1 ¥2 ¥4 ¥6 ¥10 ¥20
扫码支付:¥1
获取中
扫码支付

您的余额不足,请更换扫码支付或充值

打赏作者

实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

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

余额充值