GLONet-FLAX/src/schedulers.py

""" 
    @name schedulers.py   
    @info This module contains implementations of learning rate schedulers.
    @organization: University Laval
    @author     Gabriel Ramos
    @email      gabriel.ramos.1@ulaval.ca
"""

from typing import Callable
import ml_collections
import optax


def create_learning_rate_fn(
    config: ml_collections.ConfigDict,
    base_learning_rate: float,
    steps_per_epoch: int
    ) -> Callable:
    """Create learning rate schedule."""
    
    warmup_fn = optax.linear_schedule(
        init_value=0., 
        end_value=base_learning_rate,
        transition_steps=config.warmup_epochs * steps_per_epoch
    )
    cosine_epochs = max(
        config.num_epochs - config.warmup_epochs, 1
    )
    cosine_fn = optax.cosine_decay_schedule(
        init_value=base_learning_rate,
        decay_steps=cosine_epochs * steps_per_epoch
    )
    schedule_fn = optax.join_schedules(
        schedules=[warmup_fn, cosine_fn],
        boundaries=[config.warmup_epochs * steps_per_epoch]
    )
    return schedule_fn