torch加载与训练模型并对新模型进行初始化

https://blog.csdn.net/Jee_King/article/details/86423274

主要是根据这个博文进行操作，其中由于有些层无法更名所以利用pop把这些层从预训练模型中进行删除。

print('loading pretrained origin_model from {0}'.format("trained_model/mixed_second_finetune_acc97p7.pth"))
# 导入已经训练好的crnn模型
origin_model = torch.load("trained_model/mixed_second_finetune_acc97p7.pth")

# 打印模型信息
# for i in origin_model:
#     print(i, origin_model[i].size())

# 删除不相同的层
origin_model.pop('rnn.1.embedding.weight')
origin_model.pop('rnn.1.embedding.bias')

# 打印更新后模型信息
for i in origin_model:
    print(i, origin_model[i].size())

# 创建新模型并获取新字典
model = re_crnn.CRNN(32, 1, new_nclass, 256)
model_dict = model.state_dict()

# 打印新模型字典
# for i in model_dict:
#     print(i, model_dict[i].size())

# 初始化权重
new_state_dict = {k:v for k,v in origin_model.items() if k in model_dict}

model_dict.update(new_state_dict)
model.load_state_dict(model_dict)

# 打印权重信息观察
# for name, para in origin_model.named_parameters():
#     print(name, torch.max(para))
# for name, para in model.named_parameters():
#     print(name, torch.max(para))