06-05 03:01 阅读 175

pytorch 预训练模型读取修改相关参数的填坑问题

这篇文章主要介绍了pytorch 预训练模型读取修改相关参数的填坑问题，具有很好的参考价值，希望对大家有所帮助。如有错误或未考虑完全的地方，望不吝赐教

pytorch 预训练模型读取修改相关参数的填坑

修改部分层，仍然调用之前的模型参数。

resnet = resnet50(pretrained=False)

resnet.load_state_dict(torch.load(args.predir))

res_conv31 = Bottleneck_dilated(1024, 256,dilated_rate = 2)

print("---------------------",res_conv31)

print("---------------------",resnet.layer3[1])

res_conv31.load_state_dict(resnet.layer3[1].state_dict())

网络预训练模型与之前的模型对应不上，名称差个前缀

model_dict = model.state_dict()

# print(model_dict)

pretrained_dict = torch.load("/yzc/reid_testpcb/se_resnet50-ce0d4300.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)

最后是修改参数名拿来用的，

from collections import OrderedDict

pretrained_dict = torch.load('premodel')

new_state_dict = OrderedDict()

# for k, v in mgn_state_dict.items():

#     name = k[7:]  # remove `module.`

#     new_state_dict[name] = v

# self.model = self.model.load_state_dict(new_state_dict)

for k, v in pretrained_dict.items():

    name = "model.module."+k   # remove `module.`

    # print(name)

    new_state_dict[name] = v

self.model.load_state_dict(new_state_dict)

大家在学习pytorch时，可能想利用pytorch进行fine-tune,但是又烦恼于参数的加载问题。下面我将讲诉我的使用心得。

1

2

3

4

5

6

7

#注意：需要重新训练的层的名字要和之前的不同。

model=resnet()#自己构建的模型，以resnet为例

model_dict = model.state_dict()

pretrained_dict = torch.load('xxx.pkl')

pretrained_dict = {k: v for k, v in pretrained_dict.items() if k in model_dict}

model_dict.update(pretrained_dict)

model.load_state_dict(model_dict)

Step2：固定部分参数

1

2

3

4

5

#k是可训练参数的名字，v是包含可训练参数的一个实体

#可以先print（k）,找到自己想进行调整的层，并将该层的名字加入到if语句中：

for k,v in model.named_parameters():

    if k!='xxx.weight' and k!='xxx.bias' :

        v.requires_grad=False#固定参数

1 2	`#将要训练的参数放入优化器` `optimizer2=torch.optim.Adam(params=[model.xxx.weight，model.xxx.bias],lr=learning_rate,betas=(0.9,0.999),weight_decay=1e-5)`

debug之后，程序正常运行，最好检查一下网络的参数是否真的被固定了，如何没固定，网络的状态接近于重新训练，可能会导致网络性能不稳定，也没办法得到想要得到的性能提升。

1

2

3

for k,v in model.named_parameters():

   if k!='xxx.weight' and k!='xxx.bias' :

   print(v.requires_grad)#理想状态下，所有值都是False

需要注意的是，操作失误最大的影响是，loss函数几乎不会发生变化，一直处于最开始的状态，这很可能是因为所有参数都被固定了。

以上为个人经验，希望能给大家一个参考，也希望大家多多支持脚本之家。

您可能感兴趣的文章:

原文链接：https://blog.csdn.net/chanbo8205/article/details/89923453