Class CosineScheduler

vtp/models/utils/text_utils.py:160–207

Cosine annealing scheduler with optional warmup and freeze periods. Supports linear warmup followed by cosine annealing decay with optional initial freeze period for stable training.

    return tuple(x)

class CosineScheduler(object):
    """
    Cosine annealing scheduler with optional warmup and freeze periods.

    Supports linear warmup followed by cosine annealing decay with optional
    initial freeze period for stable training.
    """

    def __init__(self, base_value, final_value, total_iters, warmup_iters=0, start_warmup_value=0, freeze_iters=0):
        """
        Initialize cosine scheduler.

        Args:
            base_value (float): Initial value after warmup
            final_value (float): Final value at end of schedule
            total_iters (int): Total number of iterations
            warmup_iters (int): Number of warmup iterations
            start_warmup_value (float): Starting value for warmup
            freeze_iters (int): Number of initial freeze iterations
        """
        super().__init__()
        self.final_value = final_value
        self.total_iters = total_iters

        freeze_schedule = np.zeros((freeze_iters))

        warmup_schedule = np.linspace(start_warmup_value, base_value, warmup_iters)

        iters = np.arange(total_iters - warmup_iters - freeze_iters)
        schedule = final_value + 0.5 * (base_value - final_value) * (1 + np.cos(np.pi * iters / len(iters)))
        self.schedule = np.concatenate((freeze_schedule, warmup_schedule, schedule))

        assert len(self.schedule) == self.total_iters

    def __getitem__(self, it):
        """
        Get scheduled value for given iteration.

        Args:
            it (int): Current iteration

        Returns:
            float: Scheduled value for this iteration
        """
        if it >= self.total_iters:
            return self.final_value
        else:
            return self.schedule[it]