ml.lr_schedulers.linear

Defines a linear warmup and decay learning rate scheduler.

This scheduler first warms up some number of steps, then smoothly decays until the end of training.

class ml.lr_schedulers.linear.LinearLRSchedulerConfig(name: str = '???', warmup_steps: int = '???', total_steps: int = '${task.max_steps}', warmup_percent: float = 0.01, min_scale: float = 0.0001, decay: bool = True)

Bases: BaseLRSchedulerConfig

warmup_steps: int = '???'

total_steps: int = '${task.max_steps}'

warmup_percent: float = 0.01

min_scale: float = 0.0001

decay: bool = True

classmethod resolve(config: LinearLRSchedulerConfig) → None

Runs post-construction config resolution.

Parameters:

config – The config to resolve

class ml.lr_schedulers.linear.LinearLRScheduler(config: BaseConfigT)

Bases: BaseLRScheduler[LinearLRSchedulerConfig]

get_lr_scale(state: State) → float

Given a state, returns the current learning rate.

Parameters:

state – The current trainer state

Returns:

The computed learning rate to use