Transformers学习率调度器与PyTorch原生调度器的区别及自定义方法

是的，二者核心差异仅为是否包含预热阶段逻辑：

• get_constant_schedule的学习率乘数固定为1.0，全程使用优化器设置的初始学习率，无预热阶段
• get_constant_schedule_with_warmup会在warmup_steps步内将学习率从0线性提升到初始学习率，后续阶段乘数固定为1.0，和前者行为完全一致
  二者底层均返回PyTorch的LambdaLR实例，除了warmup相关逻辑外没有其他区别。

Transformers官方提供的调度器本质是对LambdaLR的封装，核心是自定义学习率乘数函数：输入为当前训练步数，输出为对应步数的学习率相对于初始学习率的乘数。要实现PyTorch原生调度器的效果，只需要把原生调度器的学习率计算逻辑转化为lambda函数即可，还可以按需叠加warmup等自定义逻辑。

通用实现模板 #

首先导入依赖：

from torch.optim import Optimizer
from torch.optim.lr_scheduler import LambdaLR

自定义调度器的基础结构如下：

def get_custom_scheduler(optimizer: Optimizer, **kwargs):
    def lr_lambda(current_step: int):
        # 自定义学习率乘数计算逻辑
        # 返回值 = 当前步学习率 / 初始学习率
        pass
    return LambdaLR(optimizer, lr_lambda)

1. 实现MultiplicativeLR效果（可选加warmup） #

MultiplicativeLR逻辑为每一步学习率乘上固定系数，叠加warmup的实现示例：

def get_multiplicative_schedule_with_warmup(
    optimizer: Optimizer,
    num_warmup_steps: int,
    multiplicative_factor: float
):
    def lr_lambda(current_step: int):
        if current_step < num_warmup_steps:
            # 预热阶段线性上升
            return float(current_step) / float(max(1, num_warmup_steps))
        # 预热结束后每步乘固定系数
        return multiplicative_factor ** (current_step - num_warmup_steps)
    return LambdaLR(optimizer, lr_lambda)

2. 实现StepLR效果（可选加warmup） #

StepLR逻辑为每经过step_size步，学习率乘衰减系数gamma：

def get_step_schedule_with_warmup(
    optimizer: Optimizer,
    num_warmup_steps: int,
    step_size: int,
    gamma: float
):
    def lr_lambda(current_step: int):
        if current_step < num_warmup_steps:
            return float(current_step) / float(max(1, num_warmup_steps))
        post_warmup_step = current_step - num_warmup_steps
        # 计算已经过的衰减周期数
        num_cycles = post_warmup_step // step_size
        return gamma ** num_cycles
    return LambdaLR(optimizer, lr_lambda)

3. 实现ExponentialLR效果（可选加warmup） #

ExponentialLR逻辑为每一步学习率乘衰减系数gamma，和固定系数的MultiplicativeLR行为一致：

def get_exponential_schedule_with_warmup(
    optimizer: Optimizer,
    num_warmup_steps: int,
    gamma: float
):
    def lr_lambda(current_step: int):
        if current_step < num_warmup_steps:
            return float(current_step) / float(max(1, num_warmup_steps))
        return gamma ** (current_step - num_warmup_steps)
    return LambdaLR(optimizer, lr_lambda)

如果不需要warmup逻辑，直接删除warmup相关的判断分支，返回对应计算逻辑的乘数值即可，和原生PyTorch调度器行为完全一致。

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