PyTorch中如何冻结其余参数仅训练MoviNet自定义分类器？

1. 先冻结模型全部参数 #

加载完成预训练模型后，先全局关闭所有参数的梯度计算：

import torch

# 假设你已经完成自定义结构的模型实例化，变量名为model
for param in model.parameters():
    param.requires_grad = False

2. 仅解锁自定义clfr分类器的参数 #

遍历参数，将前缀为clfr.的参数梯度打开：

for name, param in model.named_parameters():
    if name.startswith("clfr."):
        param.requires_grad = True

3. 迁移原有预训练分类层权重到自定义结构 #

把原有预训练模型中分类层的权重，对应赋值到修改后的clfr层对应位置：

# 加载原始MoviNet预训练权重，假设预训练权重变量为pretrain_state_dict
pretrain_state_dict = torch.load("你的预训练权重路径", map_location="cpu")

# 建立新旧层参数名映射关系
weight_map = {
    "classifier.0.conv_1.conv2d.weight": "clfr.0.multi_head.0.head2.0.conv_1.conv2d.weight",
    "classifier.0.conv_1.conv2d.bias": "clfr.0.multi_head.0.head2.0.conv_1.conv2d.bias",
    "classifier.3.conv_1.conv2d.weight": "clfr.0.multi_head.0.head1.weight",
    "classifier.3.conv_1.conv2d.bias": "clfr.0.multi_head.0.head1.bias"
}

# 替换预训练权重的key为自定义结构的key
for old_key, new_key in weight_map.items():
    if old_key in pretrain_state_dict:
        pretrain_state_dict[new_key] = pretrain_state_dict.pop(old_key)

# 加载权重到自定义模型，strict=False忽略不匹配的key
model.load_state_dict(pretrain_state_dict, strict=False)

4. 优化器仅传入可训练参数 #

初始化优化器的时候过滤掉不需要更新的冻结参数，节省显存：

optimizer = torch.optim.Adam(
    filter(lambda p: p.requires_grad, model.parameters()),
    lr=1e-4 # 分类层学习率可以适当设高一点
)

验证配置是否生效 #

执行以下代码，输出的可训练参数应该仅有你自定义的clfr层相关参数：

trainable_params = [name for name, param in model.named_parameters() if param.requires_grad]
print("可训练参数列表：")
for name in trainable_params:
    print(name)

内容的提问来源于stack exchange，提问作者Tariq Hussain