scheduler_builder

`optimus_dl.recipe.train.builders.scheduler_builder` ¶

Scheduler builder mixin for building learning rate schedulers.

`SchedulerBuilder` ¶

Builder class responsible for creating the learning rate scheduler.

Instantiates a scheduler (e.g., CosineAnnealing, WSD) and associates it with the optimizer. It ensures the scheduler is aware of the total training iterations.

Parameters:

Name	Type	Description	Default
`cfg`	`SchedulerBuilderConfig`	Builder configuration.	required
`lr_scheduler_config`	`RegistryConfig \| None`	Configuration for the scheduler itself (can be None).	required
`optimization_config`	`OptimizationConfig`	Optimization settings (needed for total iterations).	required

Source code in optimus_dl/recipe/train/builders/scheduler_builder.py

class SchedulerBuilder:
    """Builder class responsible for creating the learning rate scheduler.

    Instantiates a scheduler (e.g., CosineAnnealing, WSD) and associates it
    with the optimizer. It ensures the scheduler is aware of the total training
    iterations.

    Args:
        cfg: Builder configuration.
        lr_scheduler_config: Configuration for the scheduler itself (can be None).
        optimization_config: Optimization settings (needed for total iterations).
    """

    def __init__(
        self,
        cfg: SchedulerBuilderConfig,
        lr_scheduler_config: RegistryConfig | None,
        optimization_config: OptimizationConfig,
        **kwargs: Any,
    ):
        self.lr_scheduler_config = lr_scheduler_config
        self.optimization_config = optimization_config

    def build_lr_scheduler(self, optimizer: Optimizer, **kwargs):
        """Build and validate the learning rate scheduler.

        Args:
            optimizer: The optimizer to schedule.
            **kwargs: Additional arguments.

        Returns:
            Instantiated LR Scheduler or None if not configured.
        """
        if self.lr_scheduler_config is None:
            return None
        lr_scheduler = build(
            "lr_scheduler",
            cfg=self.lr_scheduler_config,
            optimizer=optimizer,
            iterations=self.optimization_config.iterations,
            **kwargs,
        )
        if lr_scheduler is not None:
            logger.info(f"LR Scheduler \n{lr_scheduler}")
        return lr_scheduler

`build_lr_scheduler(optimizer, **kwargs)` ¶

Build and validate the learning rate scheduler.

Parameters:

Name	Type	Description	Default
`optimizer`	`Optimizer`	The optimizer to schedule.	required
`**kwargs`		Additional arguments.	`{}`

Returns:

Type	Description
	Instantiated LR Scheduler or None if not configured.

Source code in optimus_dl/recipe/train/builders/scheduler_builder.py

def build_lr_scheduler(self, optimizer: Optimizer, **kwargs):
    """Build and validate the learning rate scheduler.

    Args:
        optimizer: The optimizer to schedule.
        **kwargs: Additional arguments.

    Returns:
        Instantiated LR Scheduler or None if not configured.
    """
    if self.lr_scheduler_config is None:
        return None
    lr_scheduler = build(
        "lr_scheduler",
        cfg=self.lr_scheduler_config,
        optimizer=optimizer,
        iterations=self.optimization_config.iterations,
        **kwargs,
    )
    if lr_scheduler is not None:
        logger.info(f"LR Scheduler \n{lr_scheduler}")
    return lr_scheduler