PyTorch-Tutorial
PyTorch-Tutorial copied to clipboard
关于模型导入预训练参数问题
自己的深度学习模型可以导入已经预训练好的参数,采用如下的代码可以导入原模型与预训练模型名字相同层的参数。
pretrained_model = models.xxx(pretrained=True) my_model = net() pretrained_dict = pretrained_model.state_dict() model_dict = my_model.state_dict() pretrained_dict = {k: v for k, v in pretrained_dict.items() if k in model_dict} model_dict.update(pretrained_dict) my_model.load_state_dict(model_dict)
但是如果我的模型和预训练好的模型结构不是完全相同,例如,原预训练模型是若干个conv_2d直接连接而成,而我的模型在每两个conv_2d中间会有一个nn.BatchNorm层,这样导致想要导入参数的层和我模型中的层有错开的情况,再使用上述代码进行导入参数时,由于名字没有办法匹配上,所以就没法按照原先想的吧conv_2d的参数互相导入,请问我说的这种情况有办法解决吗?