Skip to content

olm.train.schedulers.linear

Linear learning rate scheduler.

Classes

LinearDecayLR(*args, **kwargs) Simple linear decay scheduler that decays to zero.
LinearLR(*args, **kwargs) Linear learning rate scheduler.

class olm.train.schedulers.linear.LinearDecayLR(*args: Any, **kwargs: Any)

Bases: SchedulerBase

Simple linear decay scheduler that decays to zero.

This is a simplified version that always decays to 0 from the initial LR.

  • Parameters:
  • optimizer – Wrapped optimizer.
  • total_steps – Total number of steps to decay over.
  • last_epoch – The index of last epoch (default: -1).

Example

>>> from olm.train.schedulers import LinearDecayLR
>>> scheduler = LinearDecayLR(optimizer, total_steps=1000)
>>> for step in range(total_steps):
...     train(...)
...     scheduler.step()

get_lr()

Compute learning rate using linear decay.

class olm.train.schedulers.linear.LinearLR(*args: Any, **kwargs: Any)

Bases: SchedulerBase

Linear learning rate scheduler.

Linearly decreases (or increases) the learning rate from the initial learning rate to end_lr over total_steps.

  • Parameters:
  • optimizer – Wrapped optimizer.
  • total_steps – Total number of steps for the schedule.
  • end_lr – Target learning rate at the end (default: 0).
  • start_factor – Initial learning rate multiplier (default: 1.0).
  • last_epoch – The index of last epoch (default: -1).

Example

>>> from olm.train.schedulers import LinearLR
>>> # Decay from initial LR to 0
>>> scheduler = LinearLR(optimizer, total_steps=1000, end_lr=0)
>>> for step in range(total_steps):
...     train(...)
...     scheduler.step()

get_lr()

Compute learning rate using linear interpolation.

class olm.train.schedulers.linear.SchedulerBase(*args: Any, **kwargs: Any)

Bases: _LRScheduler, ABC

Base class for all OLM learning rate schedulers.

This class extends PyTorch’s _LRScheduler and provides a consistent interface for implementing custom learning rate schedules. All OLM schedulers should inherit from this class to maintain uniformity.

Subclasses must implement: : - get_lr(): Compute the learning rate for the current step - _get_closed_form_lr() (optional): Closed-form solution for efficiency

  • Parameters:
  • optimizer – Wrapped PyTorch optimizer.
  • last_epoch – The index of the last epoch (default: -1).
  • verbose – If True, prints a message to stdout for each update (default: False).

Example

>>> class MyScheduler(SchedulerBase):
...     def __init__(self, optimizer, param, last_epoch=-1):
...         self.param = param
...         super().__init__(optimizer, last_epoch)
...
...     def get_lr(self):
...         # Custom logic here
...         return [base_lr * self.param for base_lr in self.base_lrs]

get_last_lr() → List[float]

Return last computed learning rate by current scheduler.

  • Returns: List of last computed learning rates.

abstractmethod get_lr() → List[float]

Compute learning rate for each parameter group.

This method must be implemented by subclasses to define the learning rate schedule logic.

  • Returns: List of learning rates, one per parameter group.

load_state_dict(state_dict)

Load the scheduler state from a checkpoint.

  • Parameters: state_dict – Scheduler state returned by state_dict().

state_dict()

Returns the state of the scheduler as a dict.

Contains all non-callable attributes that are specific to the scheduler and required for checkpointing.