burn_mamba/utils/scheduler/
mod.rs

1// copied from:
2// https://github.com/huy209vn/burn-jepa/blob/588d3654fbcfdcfce2ecdb7bcaf7a2e5bd5a70ea/src/train/scheduler.rs
3// slight adaptions: added Config derives and a unified enum.
4
5//! Learning rate schedulers for controlling the optimization process.
6//!
7//! This module provides various strategies to adjust the learning rate during training,
8//! such as cosine annealing with linear warmup, to improve model convergence and performance.
9
10use burn::prelude::*;
11use std::f64::consts::PI;
12
13/// A learning-rate schedule, dispatching to one of the concrete schedulers.
14#[derive(Config, Debug)]
15pub enum Lr {
16    /// Cosine annealing with linear warmup. See [`CosineAnnealingLr`].
17    CosineAnnealing(CosineAnnealingLr),
18    /// Fixed learning rate for all steps. See [`ConstantLr`].
19    Constant(ConstantLr),
20}
21
22impl Lr {
23    /// Learning rate for the given (0-indexed) training `step`.
24    pub fn get_lr(&self, step: usize) -> f64 {
25        match self {
26            Lr::CosineAnnealing(inner) => inner.get_lr(step),
27            Lr::Constant(inner) => inner.get_lr(step),
28        }
29    }
30}
31
32/// # Cosine Annealing Learning Rate Scheduler with Linear Warmup.
33///
34/// This scheduler:
35/// 1. Linearly increases LR from 0 to `max_lr` during warmup phase
36/// 2. Applies cosine annealing from `max_lr` to `min_lr` after warmup
37///
38/// This is a common pattern in modern deep learning training.
39#[derive(Config, Debug)]
40pub struct CosineAnnealingLr {
41    /// The maximum learning rate (reached after warmup)
42    #[config(default = 1e-4)]
43    pub max_lr: f64,
44    /// The minimum learning rate (reached at end of training)
45    #[config(default = 1e-6)]
46    pub min_lr: f64,
47    /// The total number of training steps
48    pub total_steps: usize,
49    /// The number of warmup steps
50    #[config(default = 0)]
51    pub warmup_steps: usize,
52}
53
54impl CosineAnnealingLr {
55    /// Get the learning rate for the current training step.
56    ///
57    /// # Arguments
58    /// * `step` - Current training step (0-indexed)
59    ///
60    /// # Returns
61    /// * Learning rate for this step
62    pub fn get_lr(&self, step: usize) -> f64 {
63        // Warmup phase: linear increase from 0 to max_lr
64        if step < self.warmup_steps {
65            return self.max_lr * (step as f64) / (self.warmup_steps as f64);
66        }
67
68        // After total_steps, return min_lr
69        if step >= self.total_steps {
70            return self.min_lr;
71        }
72
73        // Cosine annealing phase
74        let progress =
75            (step - self.warmup_steps) as f64 / (self.total_steps - self.warmup_steps) as f64;
76        self.min_lr + 0.5 * (self.max_lr - self.min_lr) * (1.0 + (PI * progress).cos())
77    }
78}
79
80/// # Constant Learning Rate Scheduler.
81///
82/// Simply returns a fixed learning rate for all steps.
83/// Useful for simple experiments or when learning rate scheduling is not needed.
84#[derive(Config, Debug)]
85pub struct ConstantLr {
86    /// The fixed learning rate returned for every step.
87    #[config(default = 1e-4)]
88    pub lr: f64,
89}
90
91impl ConstantLr {
92    /// Returns the fixed learning rate (independent of `step`).
93    pub fn get_lr(&self, _step: usize) -> f64 {
94        self.lr
95    }
96}
97
98#[cfg(test)]
99mod tests;
burn_mamba/utils/scheduler/mod.rs

burn_mamba/utils/scheduler/
mod.rs
