修改部分层,仍然调用之前的模型参数。
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)