分层学习率:为不同层设置差异化的更新速度
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.optimizers 的 learning_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))