在PyTorch Lightning中手动配置学习率调度的问题与咨询

你想要实现的是每200个epoch将学习率减半的调度逻辑，原生PyTorch中的写法如下：

for epoch in range(0, 600):
    if (epoch + 1) % 200 == 0:
        for g in optimizer.param_groups:
            g['lr'] *= 0.5

最初代码的错误原因 #

你最初尝试在training_step中通过self.optimizers()[0]获取优化器，触发了TypeError: 'LightningAdamW' object is not subscriptable错误。这是因为self.optimizers()返回的是单个优化器实例（而非列表），所以不能用下标[0]去访问。

已验证的可行解决方法 #

你修改后的代码直接遍历self.optimizers().param_groups调整学习率是正确的，代码如下：

def training_step(self, batch, batch_idx):
    X,y = batch
    preds = self.forward(X)
    loss = self.loss(preds,y)
    self.log('train_loss', loss, prog_bar=False)
    self.train_losses.append(loss)
    
    # 每200个epoch将学习率减半（注意这里的epoch判断和原生逻辑对应，需确认是否是epoch+1触发）
    if (self.current_epoch + 1) % 200 == 0 :
        for g in self.optimizers().param_groups:
            g['lr'] *= 0.5
    
    return loss

关于是否应该放到Callback中的疑问 #

答案是更推荐将学习率调度逻辑放到Callback中，原因如下：

• 职责分离：training_step的核心职责是处理单步训练的前向传播、损失计算，将学习率调度这类全局流程控制逻辑放进去会让代码耦合度变高，可读性变差。
• 符合PL设计理念：PyTorch Lightning的Callback机制就是专门用来处理训练流程中的各类钩子事件（比如epoch开始/结束、训练开始/结束等），用Callback能让代码更模块化，后续复用或修改调度逻辑也更方便。
• 可维护性更强：把调度逻辑抽离到独立的Callback中，其他开发者能快速找到和调整学习率相关的配置，不需要在训练步骤代码中寻找零散的逻辑。

自定义Callback实现示例 #

你可以写一个简单的自定义Callback来实现这个逻辑：

from pytorch_lightning.callbacks import Callback

class StepDecayLRScheduler(Callback):
    def on_epoch_end(self, trainer, pl_module):
        # 对应原生逻辑：每200个epoch（epoch+1能被200整除）调整学习率
        if (trainer.current_epoch + 1) % 200 == 0:
            optimizer = trainer.optimizers[0]
            for param_group in optimizer.param_groups:
                param_group['lr'] *= 0.5

然后在初始化Trainer时添加这个Callback：

trainer = Trainer(callbacks=[StepDecayLRScheduler()])

这样既实现了需求，又保持了代码的整洁和模块化。

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