迁移学习在计算机视觉领域中是一种很流行的方法,因为它可以建立精确的模型,耗时更短。利用迁移学习,不是从零开始学习,而是从之前解决各种问题时学到的模式开始。这样,我们就可以利用以前的学习成果。
一:使用预训练权重
在计算机视觉领域中,迁移学习通常是通过使用预训练模型来表示的。预训练模型是在大型基准数据集上训练的模型,用于解决相似的问题。由于训练这种模型的计算成本较高,因此,导入已发布的成果并使用相应的模型是比较常见的做法。例如,在目标检测任务中,首先要利用主干神经网络进行特征提取,这里使用的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, )