declearn.optimizer.schedulers.CosineAnnealing

Bases: Scheduler

Cosine Annealing scheduler.

This scheduler implements a cosine annealing that results in the scheduled rate decreasing at each training step or round until a given index, beyond which it is constant.

It can be considered as a specific kind of decay rule that is parameterized to reach a given constant value after a given duration.

Source code in declearn/optimizer/schedulers/_cosine.py

class CosineAnnealing(Scheduler):
    """Cosine Annealing scheduler.

    This scheduler implements a cosine annealing that results
    in the scheduled rate decreasing at each training step or
    round until a given index, beyond which it is constant.

    It can be considered as a specific kind of decay rule that
    is parameterized to reach a given constant value after a
    given duration.
    """

    name = "cosine-annealing"

    def __init__(
        self,
        base: float,
        max_lr: float,
        duration: int,
        step_level: bool = True,
    ) -> None:
        """Instantiate the cosine annealing scheduler.

        Parameters
        ----------
        base:
            Minimum learning rate towards which to decrease.
        max_lr:
            Maximum learning rate, from which to start.
        duration:
            Number of steps or rounds during which to carry the cosine
            annealing. Beyond that, constantly use the `base` value.
        step_level:
            Whether to decay after each step rather than after each round.
        """
        super().__init__(base)
        self.max_lr = max_lr
        self.duration = duration
        self.step_level = step_level

    def compute_value(
        self,
        step: int,
        round_: int,
    ) -> float:
        unit = step if self.step_level else round_
        if unit >= self.duration:
            return self.base
        cosine = 1 + math.cos(unit / self.duration * math.pi)
        return self.base + 0.5 * (self.max_lr - self.base) * cosine

    def get_config(
        self,
    ) -> Dict[str, Any]:
        config = super().get_config()
        config["max_lr"] = self.max_lr
        config["duration"] = self.duration
        config["step_level"] = self.step_level
        return config

__init__(base, max_lr, duration, step_level=True)

Instantiate the cosine annealing scheduler.

Parameters:

Name	Type	Description	Default
base	float	Minimum learning rate towards which to decrease.	required
max_lr	float	Maximum learning rate, from which to start.	required
duration	int	Number of steps or rounds during which to carry the cosine annealing. Beyond that, constantly use the base value.	required
step_level	bool	Whether to decay after each step rather than after each round.	True

Source code in declearn/optimizer/schedulers/_cosine.py

def __init__(
    self,
    base: float,
    max_lr: float,
    duration: int,
    step_level: bool = True,
) -> None:
    """Instantiate the cosine annealing scheduler.

    Parameters
    ----------
    base:
        Minimum learning rate towards which to decrease.
    max_lr:
        Maximum learning rate, from which to start.
    duration:
        Number of steps or rounds during which to carry the cosine
        annealing. Beyond that, constantly use the `base` value.
    step_level:
        Whether to decay after each step rather than after each round.
    """
    super().__init__(base)
    self.max_lr = max_lr
    self.duration = duration
    self.step_level = step_level