Pytorch 加载、查看预训练模型参数、使用部分预训练模型参数初始化网络(以层为单位按需初始化)

一:查看预训练模型的参数(以Alexnet为例):

path = 'I:/迅雷下载/alexnet-owt-4df8aa71.pth'
pretrained_dict = torch.load(path)
for k, v in pretrained_dict.items():  # k 参数名 v 对应参数值
        print(k)

运行结果:                                             参数名对应AlexNet网络的位置

 二:使用部分预训练模型参数初始化网络

方式一: 自己网络和预训练网络结构一致的层,使用预训练网络对应层的参数批量初始化

model_dict = model.state_dict()                                    # 取出自己网络的参数字典
pretrained_dict = torch.load("I:/迅雷下载/alexnet-owt-4df8aa71.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)

方式二:自己网络和预训练网络结构一致的层,按层初始化

# 加粗自己定义一个网络叫CNN
model = CNN()
model_dict = model.state_dict()                                    # 取出自己网络的参数

for k, v in model_dict.items():                                    # 查看自己网络参数各层叫什么名称
       print(k)

pretrained_dict = torch.load("I:/迅雷下载/alexnet-owt-4df8aa71.pth")# 加载预训练网络的参数
for k, v in pretrained_dict.items():                                    # 查看预训练网络参数各层叫什么名称
       print(k)


# 对应层赋值初始化
model_dict['conv1.0.weight'] = pretrained_dict['features.0.weight'] # 将自己网络的conv1.0层的权重初始化为预训练网络features.0层的权重
model_dict['conv1.0.bias'] = pretrained_dict['features.0.bias']    # 将自己网络的conv1.0层的偏置项初始化为预训练网络features.0层的偏置项

model_dict['conv2.1.weight'] = pretrained_dict['features.3.weight']
model_dict['conv1.1.bias'] = pretrained_dict['features.3.bias']

model_dict['conv2.1.weight'] = pretrained_dict['features.6.weight']
model_dict['conv2.1.bias'] = pretrained_dict['features.6.bias']

... ...

参考:

   https://blog.csdn.net/chanbo8205/article/details/89923453

   https://blog.csdn.net/Code_Mart/article/details/88254444

已标记关键词 清除标记
相关推荐
©️2020 CSDN 皮肤主题: 大白 设计师:CSDN官方博客 返回首页