torch loads and trains the model and initializes the new model

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

The operation is mainly based on this blog post. Because some layers cannot be renamed, pop is used to delete these layers from the pre-training model.

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))