PyTorch学习率调度器step位置不当导致损失异常问题咨询

核心问题成因 #

你遇到的训练失效问题本质是对StepLR调度器的计数逻辑、调用层级理解错误，和PyTorch的版本规范无关。 #

两个疑问的具体解答 #

1. exp_lr_scheduler.step()的正确放置位置 #

你用的optim.lr_scheduler.StepLR(optimizer, step_size=7, gamma=0.1)是按epoch维度设计的调度器，它的内部计数逻辑是：每调用1次.step()，内部的epoch计数器加1，当计数器值达到step_size的整数倍时，就将学习率乘以gamma完成衰减。

• 这类按epoch更新的调度器，绝对不能放在batch遍历循环内部调用，只需要每跑完1个epoch（即所有batch都训练完）调用1次即可。

你的第一版代码把scheduler.step()放在batch循环里，每处理1个batch就调用1次，会导致计数器增长速度比设计值快几百倍：以MNIST数据集常规batch size配置为例，1个epoch大概有900多个batch，相当于跑完第一轮epoch时，调度器已经计数900多次，早就触发了数十次学习率衰减，初始0.003的学习率会在几十批之后就降到1e-7量级，参数几乎无法更新，损失自然一直卡在高位不动。

2. 关于scheduler.step()需放在optimizer.step()之后的规范说明 #

这个规范的真实含义是：在同一个调度周期内（对StepLR来说就是1个epoch），scheduler.step()要放在本周期所有optimizer.step()执行完成之后调用，不是要求你把它跟在每一次optimizer.step()后面。
你之前的错误写法属于把调用层级搞错了：把epoch级别的调度逻辑放到了batch级别执行，才导致学习率衰减速度完全失控，不是规范本身有问题。

另外你现在能跑通的第二版代码也存在小瑕疵：把scheduler.step()放在了epoch的最开头，相当于先更新调度器计数，再跑当前epoch的训练，在当前参数下虽然前7个epoch的学习率刚好是初始值，但存在两个隐患：

• 加载checkpoint续训时，会因为提前调用scheduler.step()导致学习率多衰减一次
• 如果后续换用ReduceLROnPlateau这类需要传入当前epoch损失值做判断的调度器，放在epoch开头会拿不到有效损失值，直接触发逻辑错误

标准正确写法 #

将scheduler.step()放在单个epoch的所有batch训练完成之后，也就是batch循环的外部、epoch循环的末尾：

def training(epoch):
    model.train()
    for batch_idx, (imgs, labels) in enumerate(train_loader):
        imgs = imgs.to(device=device)
        labels = labels.to(device=device)

        optimizer.zero_grad()
        outputs = model(imgs)
        loss = loss_fn(outputs, labels)

        loss.backward()
        optimizer.step()
    
        if (batch_idx + 1)% 100 == 0:
            print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format(
                epoch, (batch_idx + 1) * len(imgs), len(train_loader.dataset),
                100. * (batch_idx + 1) / len(train_loader), loss.data))
    # 正确位置：当前epoch所有batch的参数更新完成后，再调用调度器更新学习率
    exp_lr_scheduler.step()

补充注意点 #

• 不是所有调度器都按epoch更新：PyTorch内置的OneCycleLR、带warmup的自定义调度器属于batch级调度器，这类才需要放在batch循环里每批调用一次，使用前要先确认调度器的更新维度。
• StepLR的step_size单位是「调度器.step()的调用次数」，不是固定的epoch数，调用频率错了，衰减节奏就会完全失控。

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