现在位置: 首页 / Ai / 正文

lerobot学习率调度器

  • laumy
  • Ai
  • 12小时前
  • 19热度
  • 0评论

学习率调度器简介

是什么

学习率调度器（Learning Rate Scheduler）是深度学习训练中动态调整优化器学习率的工具（注意是在优化器的基础上动态调整学习率），通过优化收敛过程提升模型性能。

学习率（η）控制梯度更新步长，参数更新量 = -η × 梯度。过大步长导致震荡，过小则陷入局部最优或训练缓慢，固定学习率易导致训练初期震荡或后期收敛缓慢，调度器在训练期间自适应调整初期较高学习率加速收敛，中期稳定探索最优解，后期精细调优避免震荡。

以下是常见的学习率调度器

  • StepLR：每隔固定epoch将学习率乘以衰减因子（如step_size=10, gamma=0.5）。
  • ExponentialLR：每epoch按指数衰减（lr = lr × gamma）。
  • CosineAnnealingLR：按余弦曲线周期下降：η = η_min + 0.5×(η_max - η_min)×(1 + cos(π×t/T_max))。
  • OneCycleLR：分两阶段：线性升温至峰值，再退火至极小值。
  • PowerLR：基于幂律关系调整，与批次大小和token数量无关。

怎么用

import torch.optim as optim
from torch.optim.lr_scheduler import StepLR

# 1. 定义模型和优化器
model = ...  # 神经网络模型
optimizer = optim.SGD(model.parameters(), lr=0.1)  # 初始学习率0.1

# 2. 创建调度器（绑定优化器）
scheduler = StepLR(optimizer, step_size=30, gamma=0.5)  # 每30epoch衰减50

# 3. 训练循环
for epoch in range(100):

    # 前向传播 + 反向传播
    train(...)
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

    # 4. 更新学习率（通常在epoch结束时）
    scheduler.step()

    # 查看当前学习率
    print(f"Epoch {epoch}: LR={scheduler.get_last_lr()[0]:.6f}")

上面代码演示了基础的学习率调度使用示例，创建一个学习率调度器，传入的参数有优化器，主要的目的是让学习率调度器和优化器绑定，因为相当于是在优化器的基础上改进学习率调度。接下来就是在前向、反向传播更新完参数后调用scheduler.step()更新学习率，以便下一次计算使用。

from torch.optim.lr_scheduler import CosineAnnealingLR

optimizer = SGD(model.parameters(), lr=0.01)  # 基础学习率
total_epochs = 100
warmup_epochs = 10  # 预热epoch数

for epoch in range(total_epochs):
    # 1. 预热阶段：前10epoch线性增长
    if epoch < warmup_epochs:
        lr = 0.01 * (epoch / warmup_epochs)
        for param_group in optimizer.param_groups:
            param_group['lr'] = lr
    # 2. 余弦退火阶段
    else:
        scheduler = CosineAnnealingLR(optimizer, T_max=total_epochs-warmup_epochs, eta_min=1e-5)
        scheduler.step()

    # 训练代码同上...
    print(f"Epoch {epoch}: LR={optimizer.param_groups[0]['lr']:.6f}")

上面使用的是余弦退火+预热的方式。

lerobot调度器

抽象基类

@dataclass
class LRSchedulerConfig(draccus.ChoiceRegistry, abc.ABC):
    num_warmup_steps: int  # 预热步数（所有调度器的公共参数）

    @property
    def type(self) -> str:
        return self.get_choice_name(self.__class__)  # 获取子类注册名称（如 "diffuser"）

    @abc.abstractmethod
    def build(self, optimizer: Optimizer, num_training_steps: int) -> LRScheduler | None:
        """创建学习率调度器实例"""
        raise NotImplementedError

LRSchedulerConfig是学习率调度器配置的核心抽象，通过抽象接口定义+配置注册机制，实现了学习率调度策略的统一管理与灵活扩展。

其同样继承了abc.ABC标记为抽象基类，强制子类事项build抽象方法，同时继承draccus.ChoiceRegistry提供子类注册机制，通过 @LRSchedulerConfig.register_subclass("名称") 将子类与调度器类型绑定（如 "diffuser" → DiffuserSchedulerConfig），支持配置驱动的动态实例化。同时使用了@dataclass 装饰器自动生成构造函数、repr 等方法，简化调度器超参数的定义与管理。

其核心属性只有一个num_warmup_steps表示学习率预热步数。预热是深度学习训练的常见技巧（尤其在 Transformer 等模型中），将其作为基类字段可避免子类重复定义。也是几乎所有学习率调度器的基础参数，用于控制“预热阶段”（学习率从低到高线性增长的步数），避免训练初期因高学习率导致的不稳定。

其只有两个方法type和build，type是用于获取子类注册调度器类型名称进而匹配到对应子类，而build的方法是强制子类实现。

实例化子类

lerobot的学习率调度实现了3个DiffuserSchedulerConfig、VQBeTSchedulerConfig、CosineDecayWithWarmupSchedulerConfig子类，这里以DiffuserSchedulerConfig简单说明。

@LRSchedulerConfig.register_subclass("diffuser")
@dataclass
class DiffuserSchedulerConfig(LRSchedulerConfig):
    name: str = "cosine"  # 调度器类型（如 "cosine"、"linear"，来自 diffusers）
    num_warmup_steps: int | None = None  # 预热步数（可选，未指定则不预热）

    def build(self, optimizer: Optimizer, num_training_steps: int) -> LambdaLR:
        from diffusers.optimization import get_scheduler  # 复用 Diffusers 的调度器实现
        kwargs = {**asdict(self), "num_training_steps": num_training_steps, "optimizer": optimizer} # 构造调度器参数：类字段 + 训练相关参数
        return get_scheduler(**kwargs)  # 调用 diffusers API 创建调度器

diffusers.optimization.get_scheduler 是 Diffusers 库提供的调度器工厂函数，支持多种预定义策略（如 "cosine"、"linear"、"constant" 等）。

asdict(self)将 DiffuserSchedulerConfig 实例的字段（如 name="cosine"、num_warmup_steps=1000）转换为字典；而num_training_steps：总训练步数（调度器需基于此计算衰减周期）；最后的optimizer就是待绑定的优化器实例（调度器需调整其参数组的学习率）。

状态管理

（1）存储

def save_scheduler_state(scheduler: LRScheduler, save_dir: Path) -> None:
    state_dict = scheduler.state_dict()  # 获取调度器状态（如当前 step、预热步数等）
    write_json(state_dict, save_dir / SCHEDULER_STATE)  # 保存为 JSON 文件（如 "scheduler_state.json"）

该函数主要是将学习率调度器的状态写入到磁盘，为后续断点续训提供支持。state_dict = scheduler.state_dict()获取调度器的当前状态字典，包含调度器运行所需的所有动态信息。接着调用write_json写入到到文件中，文件名默认scheduler_state.json。

（2）加载

def load_scheduler_state(scheduler: LRScheduler, save_dir: Path) -> LRScheduler:
    # 从 JSON 加载状态，并按当前调度器结构适配（确保兼容性）
    state_dict = deserialize_json_into_object(save_dir / SCHEDULER_STATE, scheduler.state_dict())
    scheduler.load_state_dict(state_dict)  # 恢复状态
    return scheduler

该函数也比较简单，主要负责从磁盘加载之前保存的调度器状态（如当前训练步数、学习率调整进度等），使训练能够从断点处继续，确保学习率调整策略的连续性。

工程调用

命令参数

在lerobot中，启动训练是可以通过参数来实例化使用哪些调度器，如下。

  • --scheduler.type=diffuser指定使用diffuser调度类。
  • --scheduler.num_warmup_steps=1000指定预热步数。
  • --scheduler.num_warmup_steps=0.5指定余弦衰减周期。

创建

工程中通过 optim/factory.py 中的 make_optimizer_and_scheduler 函数统一创建优化器和调度器，并完成两者的绑定。

def make_optimizer_and_scheduler(
    cfg: TrainPipelineConfig, policy: PreTrainedPolicy
) -> tuple[Optimizer, LRScheduler | None]:
    # 1. 创建优化器（基于优化器配置，如 AdamConfig、MultiAdamConfig）
    optimizer = cfg.optimizer.build(policy.parameters())  # policy.parameters() 为模型参数

    # 2. 创建调度器（基于调度器配置，如 DiffuserSchedulerConfig、VQBeTSchedulerConfig）
    # 关键：通过调度器配置的 `build` 方法，将优化器作为参数传入，完成绑定。
    lr_scheduler = cfg.scheduler.build(optimizer, cfg.steps) if cfg.scheduler is not None else None

    return optimizer, lr_scheduler
  • 优化器创建：cfg.optimizer.build(...) 根据配置生成优化器实例（如 Adam、AdamW），并关联模型参数（policy.parameters()）。
  • 调度器绑定：cfg.scheduler.build(optimizer, cfg.steps) 调用调度器配置类（如 DiffuserSchedulerConfig）的 build 方法，将优化器实例作为参数传入，生成与该优化器绑定的调度器实例（如 LambdaLR）

更新

在训练过程中，如工程 scripts/train.py 中，调度器被集成到训练循环，通过 scheduler.step() 更新学习率通过调用scheduler.step() 更新学习率。

def train():
    # ... 初始化优化器、调度器 ...
    optimizer, lr_scheduler = make_optimizer_and_scheduler(cfg, policy)

    for step in range(start_step, cfg.steps):
        # ... 前向传播、损失计算、反向传播 ...
        optimizer.step()
        if lr_scheduler is not None:
            lr_scheduler.step()  # 调用调度器更新学习率

续训恢复

在断点续训是，通过 utils/train_utils.py 中的 save_training_state 和 load_training_state 函数处理断点续训，其中调用了 schedulers.py 的状态管理函数：

def save_training_state(step: int, optimizer: Optimizer, lr_scheduler: LRScheduler | None, save_dir: Path):
    # ... 保存优化器状态 ...
    if lr_scheduler is not None:
        save_scheduler_state(lr_scheduler, save_dir)  # 调用 schedulers.py 中的保存函数

def load_training_state(checkpoint_path: Path, optimizer: Optimizer, lr_scheduler: LRScheduler | None):
    # ... 加载优化器状态 ...
    if lr_scheduler is not None:
        lr_scheduler = load_scheduler_state(lr_scheduler, checkpoint_path)  # 调用 schedulers.py 中的加载函数
    return step, optimizer, lr_scheduler