前言
结论:1)部分加载模型参数的关键就是自定义选取训练好的模型的state_dict的键值对,然后更新到需要加载模型参数的新模型的state_dict中。 2)冻结部分参数的关键就是自定义设置需冻结的参数的requires_grad属性值为False,并在优化器中传入参数时,过滤掉requires_grad=False的这部分参数,使其不参与更新。
下文通过实例记录如何在pytorch中只加载部分参数,及冻结部分参数进行训练。
背景如下,有一个基于bert的内容分类器和基于bert的序列标注分类器,以下简称classifier与identifier。两个模型在特征抽取部分都采用相同的网络结构,就是bare_bert部分,但两种分类器的类型数量不一样,除此之外,bare_bert后续的网络结构也不同,最终需求为:“先训练classifier,然后将classifier的特征抽取部分bare_bert的参数共享给identifier,然后单独训练identifier。在训练identifier时,其共享的bare_bert的参数保持不变,只更新其层的网络参数。”
下面给了classifier与identifier的人造例子,有助于理解,classifier的定义如下:
classifier各层网络参数值为:
identifier定义如下:
identifier各层的参数值为:
可以发现,“fc1.weight, fc1.bias, fc2.weight, fc2.bias”是classifier与identifier都共有的,但具体值不相同,这里假设classifier是训练完之后的参数值,下面记录如何将共有的参数值共享给identifier。
共享部分参数
共享模型参数的核心工具就是模型的state_dict方法与load_state_dict方法。状态字典本质是python中的有序字典。
# *************** 自定义取出需要共享的参数 *******************
from collections import OrderedDict
temp = OrderedDict()
ide_state_dict = identifier.state_dict(destination=None)
for name, parameter in classifier.named_parameters():
if name in ide_state_dict:
temp[name] = parameter
# ************** 将共享的参数更新到需训练的模型中 ****************
ide_state_dict.update(temp) # 更新参数值
identifier.load_state_dict(ide_state_dict)
此时再查看identifier的参数值,可以发现“fc1.weight, fc1.bias, fc2.weight, fc2.bias”部分的参数值已经是classifier的了,注意此时的“requires_grad=True”,下面记录如何冻结部分参数。
冻结部分参数
模型中的Parameter本质是Tensor的子类,因此其有Tensor的所有属性,其中“requires_grad”属性决定是否需要计算梯度,默认情况下,模型网络中的参数是需要记录梯度的,因此需要将“requires_grad”设置为False,并且在优化器中过滤掉这部分参数。如下所示:
# 自定义冻结部分参数
for name, parameter in identifier.named_parameters():
if 'ide_only' not in name:
parameter.requries_grad = False
# 过滤传入优化器的参数
optimizer = torch.optim.Adam(filter(lambda p: p.requires_grad, identifier.parameters()))
参考资料
How the pytorch freeze network in some layers, only the rest of the training?