PyTorch中如何为不同参数组配置不同的学习率调度器？

要实现不同参数组使用不同的指数衰减率（gamma），不能直接用ExponentialLR（它会对所有参数组应用相同的gamma），可以用PyTorch提供的MultiplicativeLR调度器，它支持为每个参数组自定义衰减因子。

核心思路 #

MultiplicativeLR允许传入一个lambda函数，该函数接收参数组的索引或参数组本身作为输入，返回对应的学习率乘法因子（即你需要的gamma值），以此实现不同参数组的独立衰减。

完整实现代码 #

import torch
import torch.nn as nn
from torch.optim.lr_scheduler import MultiplicativeLR

# 定义模型
class NN1(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc = nn.Linear(10, 1)

class NN2(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc = nn.Linear(10, 1)

# 初始化模型和优化器
module1 = NN1()
module2 = NN2()

criterion = nn.MSELoss()
optimizer = torch.optim.Adam([ 
                {'params': module1.parameters(), 'lr':0.2},
                {'params': module2.parameters(), 'lr': 0.1}
            ])

# 定义调度器：为不同参数组设置不同gamma
# 通过参数组索引区分，0对应module1，1对应module2
scheduler = MultiplicativeLR(
    optimizer,
    lr_lambda=lambda idx: 0.8 if idx == 0 else 0.9,
    verbose=True
)

# 训练循环示例
for epoch in range(5):
    # 模拟训练流程
    inputs = torch.randn(32, 10)
    target1 = torch.randn(32, 1)
    target2 = torch.randn(32, 1)
    
    optimizer.zero_grad()
    output1 = module1(inputs)
    output2 = module2(inputs)
    loss = criterion(output1, target1) + criterion(output2, target2)
    loss.backward()
    optimizer.step()
    
    # 更新学习率
    scheduler.step()
    print(f"Epoch {epoch+1}, 当前学习率: {[round(group['lr'], 6) for group in optimizer.param_groups]}")

更灵活的写法（不依赖参数组顺序） #

如果担心参数组顺序变化导致错误，可以在定义优化器时给每个参数组添加自定义gamma键，再通过lambda函数读取该值：

# 优化器定义时添加gamma标识
optimizer = torch.optim.Adam([ 
                {'params': module1.parameters(), 'lr':0.2, 'gamma':0.8},
                {'params': module2.parameters(), 'lr': 0.1, 'gamma':0.9}
            ])

# 调度器直接读取参数组的gamma值
scheduler = MultiplicativeLR(
    optimizer,
    lr_lambda=lambda group: group['gamma'],
    verbose=True
)

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