优化器参数调度器

此 API 用于计算优化器的学习率和权重衰减。

学习率衰减和权重衰减增加函数。

class core.optimizer_param_scheduler.OptimizerParamScheduler(optimizer: megatron.core.optimizer.MegatronOptimizer, init_lr: float, max_lr: float, min_lr: float, lr_warmup_steps: int, lr_decay_steps: int, lr_decay_style: str, start_wd: float, end_wd: float, wd_incr_steps: int, wd_incr_style: str, use_checkpoint_opt_param_scheduler: Optional[bool] = True, override_opt_param_scheduler: Optional[bool] = False, wsd_decay_steps: Optional[int] = None, lr_wsd_decay_style: Optional[str] = None)

基类：object

退火学习率和权重衰减

参数

• optimizer (MegatronOptimizer) – 要使用的优化器

• init_lr (float) – 初始学习率

• max_lr (float) – 最大学习率

• min_lr (float) – 最小学习率

• lr_warmup_steps (int) – 预热步数

• lr_decay_steps (int) – 衰减步数

• lr_decay_style (str) – 学习率的衰减风格

• start_wd (float) – 初始权重衰减

• end_wd (float) – 最终权重衰减

• wd_incr_steps (int) – 权重衰减增加步数

• wd_incr_style (str) – 权重衰减增加风格

• use_checkpoint_opt_param_scheduler (bool, optional) – 是否使用检查点值作为优化器参数调度器

• override_opt_param_scheduler (bool, optional) – 是否使用类值覆盖优化器参数调度器的值

• wsd_decay_steps (int, optional) – 权重衰减的衰减步数

• lr_wsd_decay_style (str, optional) – 权重衰减衰减步数期间学习率的衰减风格

get_lr(param_group: dict) → float

学习率衰减函数来自： https://openreview.net/pdf?id=BJYwwY9ll 第 4 页

参数

param_group (dict) – 来自优化器的参数组。

get_wd() → float

权重衰减增加函数

load_state_dict(state_dict: dict) → None

加载状态字典。

参数

state_dict (dict) – 要加载的状态字典

state_dict() → dict

返回状态字典。

step(increment: int) → None

为所有参数组设置学习率。

参数

increment (int) – 要增加的步数