分层学习率:为不同层设置差异化的更新速度

FreeGuideOnline 最新 2026-06-21

python import torch import torch.optim as optim from torchvision import models

加载预训练模型

model = models.resnet18(pretrained=True)

假定我们要做10分类任务,替换最后的全连接层

num_ftrs = model.fc.in_features model.fc = torch.nn.Linear(num_ftrs, 10)

将参数按层划分成不同的组,并设置各自的学习率

params_to_update = [] for name, param in model.named_parameters(): if not param.requires_grad: continue if 'layer4' in name or 'fc' in name:
# 最高层:layer4和全连接层使用较大学习率 params_to_update.append({'params': param, 'lr': 1e-3}) elif 'layer3' in name: params_to_update.append({'params': param, 'lr': 5e-4}) elif 'layer2' in name: params_to_update.append({'params': param, 'lr': 2e-4}) else: # layer1 及更底层使用最小的学习率 params_to_update.append({'params': param, 'lr': 1e-4})

也可直接为不同组设置不同参数

optimizer = optim.SGD(params_to_update, momentum=0.9, weight_decay=1e-4)


**关键步骤解释:**
- 遍历模型的 `named_parameters()`,根据层的名字判断所属模块。
- 为每个参数创建一个字典,其中包含 `params` 和 `lr`(还可以加入 `weight_decay` 等)。
- 将这个字典列表作为优化器的输入。优化器内部会自动管理这些组的学习率。

若使用 Adam/AdamW,原理完全相同:

```python
optimizer = optim.AdamW([
    {'params': model.layer4.parameters(), 'lr': 1e-3},
    {'params': model.layer3.parameters(), 'lr': 5e-4},
    {'params': model.layer2.parameters(), 'lr': 2e-4},
    {'params': model.layer1.parameters(), 'lr': 1e-4},
    {'params': model.fc.parameters(), 'lr': 1e-3}
], weight_decay=0.01)

TensorFlow / Keras 实现

在 Keras 中,可以使用 tf.keras.optimizerslearning_rate 参数为不同的层传入不同的调度或常数,但更灵活的方式是自定义训练循环。以下展示在自定义循环中为不同层设置不同学习率:

import tensorflow as tf

model = ...  # 你的Keras模型

# 将变量按层分组,并分配学习率
var_groups = {
    'low_lr': model.layers[:5],   # 前5层
    'mid_lr': model.layers[5:10], # 中间层
    'high_lr': model.layers[10:]  # 高层及分类头
}

optimizers_and_vars = [
    (tf.keras.optimizers.Adam(1e-4), [v for l in var_groups['low_lr'] for v in l.trainable_variables]),
    (tf.keras.optimizers.Adam(5e-4), [v for l in var_groups['mid_lr'] for v in l.trainable_variables]),
    (tf.keras.optimizers.Adam(1e-3), [v for l in var_groups['high_lr'] for v in l.trainable_variables])
]

@tf.function
def train_step(x, y):
    with tf.GradientTape() as tape:
        pred = model(x, training=True)
        loss = loss_fn(y, pred)
    grads = tape.gradient(loss, model.trainable_variables)
    # 分别对各组变量应用各自的优化器
    for optimizer, var_list in optimizers_and_vars:
        optimizer.apply_gradients(zip(grads, var_list))