【Pytorch】利用预训练模型初始化backbone时的常见问题及方案

在训练模型时，我们常常需要利用预训练的baseline模型对所设计网络的backbone或部分layer进行初始化，给网络训练提供一个较好的起点，同时减少训练的时间成本。比较常见的就是利用imagenet上训练好的标准网络来初始化新网络的部分层。

在进行初始化时，往往会出现两种情况：一种是待初始化的层键值和预训练模型是匹配的；一种是二者键值不匹配，在键值名称上有少许差异。

针对这两种情况，处理方案如下：

1）新模型和baseline模型的键值匹配

#分别取出checkpoint和new model的键值对
checkpoint_dict = checkpoint.state_dict()
model_dict = model.state_dict()
#从checkpoint中找出二者相同的键值对，替换掉new model中相应的键值对，并更新网络     
backbone_dict = {k:v for k,v in checkpoint_dict.items() if k in model_dict}
model_dict.update(backbone_dict)
model.load_state_dict(model_dict)

2）新模型和baseline模型的键值名称上存在差异，比如由module.layers.0.conv1.weight 改为 modulelist.layers.0.conv1.weight，这时需要对checkpoint键值进行更新，使其与新模型键值匹配后再进行初始化

#分别取出checkpoint和new model的键值对,并更新checkpoint中键值名称
checkpoint_dict = checkpoint.state_dict()
model_dict = model.state_dict() 
new_dict = {k.replace('module', 'modulelist'):v for k,v in checkpoint_dict.items()}

#从checkpoint中找出二者相同的键值对，替换掉new model中相应的键值对，并更新网络     
backbone_dict = {k:v for k,v in checkpoint_dict.items() if k in model_dict.items()}
model_dict.update(backbone_dict)
model.load_state_dict(model_dict)

 

 

References:

[1]http://xcx1024.com/ArtInfo/839249.html （此文总结较全，推荐！）

[2]https://blog.csdn.net/Gavinmiaoc/article/details/80514528

 

注：本人经验尚浅，若笔记中存在理解错误之处，望大家批评指正。