torch权重冻结&解冻

发表于 分类于

记录pytorch中如何固定部分参数的权重，实现模型部分训练。

参数能被更新的条件：

  1. 参数允许保存梯度（即requires_grad=True）
  2. 有优化器对参数进行修改（即绑定了对应的Optimizer对象）

二者必须同时满足。

冻结

在初始化模型之后，对想要冻结的参数设置其requires_grad=False即可，即不记录梯度信息。

1
2
3
4
5
6
7
8
9
10
11
from collections.abc import Iterable

# 冻结某个某块的某几个子模块
def set_freeze_by_names(model, layer_names, freeze=True):
    if not isinstance(layer_names, Iterable):
        layer_names = [layer_names]
    for name, child in model.named_children(): # 遍历所有直接被包含的子模块
        if name not in layer_names:
            continue
        for param in child.parameters(): # 设置需要的子模块参数是否保存梯度
            param.requires_grad = not freeze

实验结果如下：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
# 假设有如下一个模型
model = nn.Sequential(
	nn.Linear(32,64),
    nn.Linear(64,128),
    nn.Linear(128,256)
)
print(list(model.named_children()))
# >>> [('0', Linear(in_features=32, out_features=64, bias=True)),
#     ('1', Linear(in_features=64, out_features=128, bias=True)),
#     ('2', Linear(in_features=128, out_features=256, bias=True))]
# 共三层子模块，名称为0,1,2

# 灵活控制冻结的子模块
set_freeze_by_names(model, ('0', '1'), freeze=True) # 设置0,1两层参数冻结
print(model[0].weight.requires_grad)
# >>> False
print(model[1].weight.requires_grad)
# >>> False
print(model[2].weight.requires_grad)
# >>> True

# 也可以通过设置in-place方法，直接冻结整个模块的所有参数
model.requires_grad_(False)

然后在训练前绑定优化器的时候只给保存梯度的参数绑定优化器即可：

1
optim = torch.optim.Adam([p for p in model.parameters() if p.requires_grad], lr=0.0001)

解冻

解冻参数遵守两步：

  1. 允许参数保存梯度，即requires_grad=True
  2. 在优化器中加入解冻的参数，使其能被更新
1
2
3
4
5
6
7
8
9
10
11
# # 灵活控制被解冻的子模块
set_freeze_by_names(model, ('0'), freeze=False) # 解冻0层
print(model[0].weight.requires_grad)
# >>> True
print(model[1].weight.requires_grad)
# >>> False
print(model[2].weight.requires_grad)
# >>> True

# 也可设置in-place方法，解冻整个模块的所有参数
model.requires_grad_(True)

然后在优化器中加入解冻出来的参数：

1
optim.add_param_group({'params': model[0].parameters()})