burn_optim/lr_scheduler/
exponential.rs

1use burn_core as burn;
2
3use super::{LrScheduler, LrSchedulerRecord, String};
4use crate::LearningRate;
5use crate::RecordState;
6use crate::lr_scheduler::module_lr_scheduler::ModuleLrScheduler;
7use burn::config::Config;
8
9/// The configuration for creating an [exponential learning rate scheduler](ExponentialLrScheduler).
10///
11/// This scheduler returns the learning rate `initial_lr` at the first step, then multiplies it by
12/// a constant `gamma` at every iteration. At any iteration `i` (which starts from 0), the learning
13/// rate is given by `initial_lr * gamma^i`.
14#[derive(Config, Debug)]
15pub struct ExponentialLrSchedulerConfig {
16    // The initial learning rate.
17    initial_lr: LearningRate,
18    // The constant that the learning rate is multiplied by on each iteration.
19    gamma: f64,
20}
21
22impl ExponentialLrSchedulerConfig {
23    /// Initializes a [exponential learning rate scheduler](ExponentialLrScheduler).
24    pub(crate) fn build(&self) -> Result<ExponentialLrScheduler, String> {
25        if self.initial_lr <= 0. || self.initial_lr > 1. {
26            return Err("Initial learning rate must be greater than 0 and at most 1".into());
27        }
28        if self.gamma <= 0. || self.gamma > 1. {
29            return Err("Gamma must be greater than 0 and at most 1".into());
30        }
31
32        Ok(ExponentialLrScheduler {
33            // Such an initial value eliminates the need for special-case handling of the first
34            // learning rate.
35            previous_lr: self.initial_lr / self.gamma,
36            gamma: self.gamma,
37        })
38    }
39
40    /// Initializes a [module learning rate scheduler](ModuleLrScheduler).
41    ///
42    /// # Errors
43    ///
44    /// An error will be returned if any of the following conditions is true:
45    ///
46    /// * `initial_lr` is out of range (0.0, 1.0]
47    /// * `gamma` is out of range (0.0, 1.0]
48    pub fn init(&self) -> Result<ModuleLrScheduler, String> {
49        self.build().map(|s| s.into())
50    }
51}
52
53/// A exponential learning rate scheduler.
54///
55/// See [ExponentialLrSchedulerConfig] for more information.
56#[derive(Clone, Copy, Debug)]
57pub struct ExponentialLrScheduler {
58    // The previous iteration's learning rate.
59    previous_lr: LearningRate,
60    // The constant that the learning rate is multiplied by on each iteration.
61    gamma: f64,
62}
63
64impl LrScheduler for ExponentialLrScheduler {
65    fn step(&mut self) -> LearningRate {
66        self.previous_lr *= self.gamma;
67        self.previous_lr
68    }
69
70    fn to_record(&self) -> LrSchedulerRecord {
71        LrSchedulerRecord::from_state(&ExponentialLrSchedulerState {
72            previous_lr: self.previous_lr,
73        })
74    }
75
76    fn load_record(&mut self, record: LrSchedulerRecord) {
77        if let Some(state) = record.into_state::<ExponentialLrSchedulerState>() {
78            self.previous_lr = state.previous_lr;
79        }
80    }
81}
82
83/// The serializable state of an [exponential scheduler](ExponentialLrScheduler).
84#[derive(RecordState, Clone, Debug)]
85pub struct ExponentialLrSchedulerState {
86    // `f64` (not the `LearningRate` alias) so the derive recognizes it as a scalar leaf.
87    previous_lr: f64,
88}
89
90#[cfg(test)]
91mod tests {
92    use super::super::test_utils;
93    use super::*;
94
95    #[test]
96    fn config_initial_lr_too_low() {
97        let r = ExponentialLrSchedulerConfig::new(0., 0.5).build();
98        assert!(r.is_err(), "Should return an error");
99        assert_eq!(
100            r.unwrap_err(),
101            "Initial learning rate must be greater than 0 and at most 1",
102            "Error messages should match",
103        );
104    }
105
106    #[test]
107    fn config_initial_lr_too_high() {
108        let r = ExponentialLrSchedulerConfig::new(1.5, 0.5).build();
109        assert!(r.is_err(), "Should return an error");
110        assert_eq!(
111            r.unwrap_err(),
112            "Initial learning rate must be greater than 0 and at most 1",
113            "Error messages should match",
114        );
115    }
116
117    #[test]
118    fn config_gamma_too_low() {
119        let r = ExponentialLrSchedulerConfig::new(0.5, 0.0).build();
120        assert!(r.is_err(), "Should return an error");
121        assert_eq!(
122            r.unwrap_err(),
123            "Gamma must be greater than 0 and at most 1",
124            "Error messages should match",
125        );
126    }
127
128    #[test]
129    fn config_gamma_too_high() {
130        let r = ExponentialLrSchedulerConfig::new(0.5, 1.5).build();
131        assert!(r.is_err(), "Should return an error");
132        assert_eq!(
133            r.unwrap_err(),
134            "Gamma must be greater than 0 and at most 1",
135            "Error messages should match",
136        );
137    }
138
139    #[test]
140    fn test_lr_change() {
141        let scheduler = ExponentialLrSchedulerConfig::new(0.8, 0.1).build().unwrap();
142        let expected_lrs = [0.8, 0.08, 0.008, 0.0008, 0.00008];
143        test_utils::check_lr_sequence(scheduler, expected_lrs);
144    }
145
146    #[test]
147    fn test_save_and_load() {
148        let scheduler = ExponentialLrSchedulerConfig::new(0.083, 0.3)
149            .build()
150            .unwrap();
151        test_utils::check_save_load(scheduler, 7);
152    }
153}
burn_optim/lr_scheduler/exponential.rs

burn_optim/lr_scheduler/
exponential.rs
