PyTorch：如何结合学习率调度器实现分层学习率衰减？

我现在在用PyTorch的OneCycleLR调度器，代码如下：

model = Model()
optim = optim.Adam(lr=0.1)
scheduler = optim.lr_scheduler.OneCycleLR(optim, max_lr=0.1)

当前学习率会在前30%训练轮次上升到0.1，之后逐渐衰减。我想在这个基础上加上分层学习率衰减：每一步先获取调度器给出的当前学习率，再给每层按0.9^(idx+1)的系数衰减，大致逻辑是：

for i in range(steps):
    lr = scheduler.get_last_lr()
    for idx, layer in enumerate(model.layers()):
        layer['lr'] = lr * 0.9 ** (idx+1)
    output = model(input)
    ...

但我有几个疑问：

• 这种方式需要重新把model.parameters()传入优化器吗？
• 此时学习率具体会怎么计算？
• 有没有更优的实现方式？
• 如何适配超大模型，避免逐个指定每层学习率的繁琐操作？

核心思路：利用优化器参数组的lr_lambda实现动态分层衰减 #

PyTorch优化器的参数组支持lr_lambda参数，它可以接收一个函数，该函数会以调度器调整后的基础学习率为输入，输出该参数组的最终学习率。这样就能完美结合调度器的动态调整和分层衰减逻辑，不需要手动循环修改每层学习率。

具体实现步骤 #

1. 按层分组模型参数：把模型的每层参数单独打包成参数组，或者按模块层级批量分组（适合超大模型）。
2. 为每个参数组设置lr_lambda：lambda函数里实现0.9^(idx+1)的衰减逻辑，其中idx是层的索引。
3. 正常使用调度器：调度器会先调整基础学习率，再通过lr_lambda计算每层的最终学习率。

代码示例 #

基础版（适用于明确分层的模型）

import torch
import torch.nn as nn
import torch.optim as optim

# 示例模型：简单的多层MLP
class Model(nn.Module):
    def __init__(self):
        super().__init__()
        self.layers = nn.ModuleList([
            nn.Linear(10, 20),
            nn.ReLU(),
            nn.Linear(20, 30),
            nn.ReLU(),
            nn.Linear(30, 1)
        ])
    
    def forward(self, x):
        for layer in self.layers:
            x = layer(x)
        return x

model = Model()

# 1. 按层创建参数组，每层对应一个lambda函数
param_groups = []
for idx, layer in enumerate(model.layers):
    # 只针对可训练参数（比如跳过ReLU的参数）
    if list(layer.parameters()):
        param_groups.append({
            'params': layer.parameters(),
            'lr_lambda': lambda epoch, idx=idx: 0.9 ** (idx + 1)
        })

# 2. 初始化优化器，基础lr设为调度器的初始值
optim = optim.Adam(param_groups, lr=0.1)

# 3. 初始化OneCycleLR调度器
scheduler = optim.lr_scheduler.OneCycleLR(optim, max_lr=0.1, total_steps=1000)

# 训练循环
for step in range(1000):
    input = torch.randn(32, 10)
    output = model(input)
    loss = output.mean()
    loss.backward()
    optim.step()
    # 调度器每步更新（OneCycleLR需要每步调用）
    scheduler.step()
    
    # 查看每层当前学习率（验证用）
    if step % 100 == 0:
        print(f"Step {step}:")
        for idx, group in enumerate(optim.param_groups):
            current_lr = group['lr']
            print(f"Layer {idx+1}: lr = {current_lr:.6f}")

超大模型适配版（自动按模块层级分组）

如果模型是超大Transformer类模型（比如BERT、GPT），可以按nn.Module的层级自动递归分组，不需要手动枚举每层：

def get_layer_param_groups(model, decay_factor=0.9):
    param_groups = []
    # 递归遍历模型的所有子模块
    def traverse(module, depth=0):
        # 只处理叶子模块（没有子模块的模块）
        if len(list(module.children())) == 0 and list(module.parameters()):
            param_groups.append({
                'params': module.parameters(),
                'lr_lambda': lambda epoch, d=depth: decay_factor ** (d + 1)
            })
        else:
            for child in module.children():
                traverse(child, depth + 1)
    traverse(model)
    return param_groups

# 假设model是一个超大Transformer模型
param_groups = get_layer_param_groups(model, decay_factor=0.9)
optim = optim.Adam(param_groups, lr=0.1)
scheduler = optim.lr_scheduler.OneCycleLR(optim, max_lr=0.1, total_steps=1000)

疑问解答 #

1. 不需要重新传入model.parameters()：优化器初始化时已经接收了按层分组的参数组，后续调度器和lr_lambda会自动处理学习率更新，不需要再手动传入参数。
2. 学习率计算逻辑：最终每层的学习率 = 调度器调整后的基础学习率 × 0.9^(idx+1)。比如调度器当前输出的基础lr是0.1，第1层的lr是0.1×0.9^1=0.09，第2层是0.1×0.9^2=0.081，以此类推，且这个值会随着调度器的每步更新动态变化。
3. 更优实现方式：上面的lr_lambda方案是PyTorch原生支持的最优方式，不需要手动循环修改学习率，避免了手动操作可能带来的错误（比如忘记更新、参数组不匹配）。
4. 超大模型适配：通过递归遍历模型的子模块自动分组，不需要手动枚举每层，代码简洁且适配任意规模的模型。

内容的提问来源于stack exchange，提问作者Long Luu