ml.lr_schedulers.base

Defines the base class and config for all learning rate schedulers.

Learning rate schedulers can be plotted using the tool in ml.tools.plot_lr_schedulers. For example:

python -m ml.lr_schedulers.scripts.plot linear /path/to/save.png

class ml.lr_schedulers.base.SchedulerAdapter(scheduler: BaseLRScheduler, optimizer: Optimizer)

Bases: object

Defines a general-purpose learning rate scheduler adapter.

last_state: State | None

state_dict() → dict[str, Any]

load_state_dict(state_dict: dict[str, Any]) → None

step(state: State) → None

class ml.lr_schedulers.base.BaseLRSchedulerConfig(name: str = '???')

Bases: BaseConfig

Defines the base config for all learning rate schedulers.

class ml.lr_schedulers.base.BaseLRScheduler(config: BaseConfigT)

Bases: BaseObject[LRSchedulerConfigT], Generic[LRSchedulerConfigT], ABC

Defines the base learning rate scheduler.

abstract 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

get(optimizer: Optimizer) → SchedulerAdapter