加载预训练权重要很久吗（迁移学习-如何使用预训练权重，冻结部分层权重训练）

        迁移学习在计算机视觉领域中是一种很流行的方法  
    
    
   ，因为它可以建立精确的模型

     


     


     

，耗时更短  
    
    
   。利用迁移学习 




 




 




，不是从零开始学习  

   

   

  ，而是从之前解决各种问题时学到的模式开始

     


     


     

。这样

     

     

     
，我们就可以利用以前的学习成果 




 




 




。

一：使用预训练权重

        在计算机视觉领域中  




  




  




，迁移学习通常是通过使用预训练模型来表示的  

   

   

  。预训练模型是在大型基准数据集上训练的模型  


  


  


 ，用于解决相似的问题

     

     

     
。由于训练这种模型的计算成本较高
     
     
     ，因此
   




   




   




，导入已发布的成果并使用相应的模型是比较常见的做法  




  




  




。例如 



 



 



，在目标检测任务中                ，首先要利用主干神经网络进行特征提取

    




    




    



，这里使用的backbone一般就是VGG    
    
    
、ResNet等神经网络














，因此在训练一个目标检测模型时                ，可以使用这些神经网络的预训练权重来将backbone的参数初始化

     



     



     


，这样在一开始就能提取到比较有效的特征  


  


  


 。

        PyTorch提供了state_dict()和load_state_dict()两个方法用来保存和加载模型参数














，前者将模型参数保存为字典形式  
    
    
   ，后者将字典形式的模型参数载入到模型当中
     
     
     。

示例：

我们先定义一个模型（省略）

model = regress()

获得当前模型的参数

model_dict = model.state_dict()

获取预训练的模型

pre_model = torch.load(premodel_path, map_location="cpu")

这里要注意

     


     


     

，torch.load()得到的结果并不一定就是模型 




 




 




，也可以是模型参数或者其他储存模型信息的字典  

   

   

  ，这里一定要据情况而论
   




   




   




。

# 就比如这里定义一个字典 s = {"state_dict":model_dict, "test1":1, "test2":2} # 先保存 torch.save(s, "mytest") # 再用load方法读取 test = torch.load("mytest", map_location="cpu") # 返回的就是个字典 print(type(test)) # <class dict>

所以一般来说要看读取的文件是什么

     

     

     
，如果以.pth结尾就是模型 



 



 



。如果是其他形式那么可以用字典映射来获取想要的信息  




  




  




，比如这里我们获取模型参数                。

test = torch.load("mytest", map_location="cpu")["state_dict"] print(type(test)) # <class collections.OrderedDict>

获取完预训练模型参数后  


  


  


 ，如果要将这个模型参数加载进我们的模型要使用load_state_dict()方法

    




    




    



。

而要注意load_state_dict()方法要求加载的模型参数键值类型和当前模型完全一致
     
     
     ，所以在载入的过程中需要匹配














。

在我们得到当前模型参数model_dict和预训练模型参数pre_model.state_dict()后
   




   




   




，

# 将模型中匹配的层暂存到一个临时的字典 # 要保证k这个层在原模型中 



 



 



，且模型结构shape要一样 temp = {k: v for k, v in pre_model.state_dict().items() if k in model_dict and np.shape(model_dict[k]) == np.shape(v)}

 接下来将temp更新到我们的模型参数里

# 这里我在debug的时候发现如果不把model_dict的数据类型改为dict会发生字典键值对无法更新的情况                ，我觉得可能是模型参数的数据类型是OrderedDict,和temp不一样 model_dict = dict(model_dict) # 更新模型参数字典 model_dict.update(temp) # 将模型参数加载到当前模型里 model.load_state_dict(model_dict, strict=True) # 这里注意strict参数为True(默认值)表示待加载模型参数和当前模型的结构完全一样

    




    




    



，如果不一样那么就要将strict设置为False # 比如我直接将temp加载到model里 model.load_state_dict(temp, strict=False)

二   


     


     


   、冻结部分层权重训练

冻结训练很简单














，只需要调整对应层的requires_grad属性就行了

这里将当前模型除了最后的全连接层外所有层权重冻结                ，仅训练最后的全连接层

for par in model.parameters(): par.requires_grad = False model.fc1.weight.requires_grad = True # 这里还要注意要把优化器调整一下

     



     



     


，只保留模型中可训练的层                。 # 可以用一个列表储存所有待训练层 train_layer = [p for p in model.parameters() if p.requires_grad == True] optimizer = torch.optim.Adam(train_layer, )