burn_optim/lr_scheduler/
step.rs

1use burn_core as burn;
2
3use burn::config::Config;
4
5use super::{LrScheduler, LrSchedulerRecord, String};
6use crate::lr_scheduler::module_lr_scheduler::ModuleLrScheduler;
7use crate::{LearningRate, RecordState};
8
9/// The configuration for create a [step learning rate scheduler](StepLrScheduler).
10///
11/// This scheduler returns the learning rate `initial_lr` from the start, and keeps doing so until
12/// the same value has been given for `step_size` times. Then it multiplies the learning rate by
13/// `gamma` before repeating the process.
14///
15/// Gamma values out of range (0.0, 1.0) and non-positive initial learning rates are acceptable, but
16/// a warning log will be output for such a value in case of mistyping.
17///
18/// ## Notes
19///
20/// The [step](StepLrScheduler::step) method of the scheduler panics if it is called more than
21/// `i32::MAX + 1` times.
22#[derive(Config, Debug)]
23pub struct StepLrSchedulerConfig {
24    // The learning rate at the initial step.
25    initial_lr: LearningRate,
26    // The number of iterations over which the learning rate remains unchanged before the next
27    // update.
28    step_size: usize,
29    /// The factor by which the learning rate is multiplied with each update. Default: 0.1.
30    #[config(default = 0.1)]
31    gamma: f64,
32}
33
34impl StepLrSchedulerConfig {
35    /// Initializes a [step learning rate scheduler](StepLrScheduler).
36    pub(crate) fn build(&self) -> Result<StepLrScheduler, String> {
37        if self.step_size == 0 {
38            return Err("Step size must be greater than 0".into());
39        }
40
41        // Atypical values of `initial_lr` and `gamma` are not rejected because they might be useful
42        // in some cases like debugging (e.g., https://datascience.stackexchange.com/q/89518).
43        if self.initial_lr <= 0.0 {
44            log::warn!(
45                "Initial learning rate value of {} is not a positive number. Ignore this warning \
46                 if it is intended.",
47                self.initial_lr
48            );
49        }
50        if self.gamma <= 0.0 || self.gamma >= 1.0 {
51            log::warn!(
52                "Gamma value of {} is out of range (0.0, 1.0). Ignore this warning if it is \
53                 intended.",
54                self.gamma
55            );
56        }
57
58        Ok(StepLrScheduler {
59            init_lr: self.initial_lr,
60            step_size: self.step_size,
61            gamma: self.gamma,
62            iter_idx: -1,
63        })
64    }
65
66    /// Initializes a [module learning rate scheduler](ModuleLrScheduler).
67    ///
68    /// # Errors
69    ///
70    /// An error will be returned if `step_size` is 0.
71    pub fn init(&self) -> Result<ModuleLrScheduler, String> {
72        self.build().map(|s| s.into())
73    }
74}
75
76/// Step learning rate scheduler.
77#[derive(Clone, Debug)]
78pub struct StepLrScheduler {
79    init_lr: LearningRate,
80    step_size: usize,
81    gamma: f64,
82    // The index of the current iteration.
83    // `i32` is used for avoiding truncating the exponent when taking powers of `gamma`.
84    iter_idx: i32,
85}
86
87impl LrScheduler for StepLrScheduler {
88    fn step(&mut self) -> LearningRate {
89        self.iter_idx = self
90            .iter_idx
91            .checked_add(1)
92            .expect("`.step()` should be called no more than `i32::MAX + 1` times");
93        // Type casting below causes no truncation, as all the values fall within the ranges.
94        self.init_lr
95            * self
96                .gamma
97                .powi((self.iter_idx as usize / self.step_size) as i32)
98    }
99
100    fn to_record(&self) -> LrSchedulerRecord {
101        LrSchedulerRecord::from_state(&StepLrSchedulerState {
102            iter_idx: self.iter_idx,
103        })
104    }
105
106    fn load_record(&mut self, record: LrSchedulerRecord) {
107        if let Some(state) = record.into_state::<StepLrSchedulerState>() {
108            self.iter_idx = state.iter_idx;
109        }
110    }
111}
112
113/// The serializable state of a [step learning rate scheduler](StepLrScheduler).
114#[derive(RecordState, Clone, Debug)]
115pub struct StepLrSchedulerState {
116    iter_idx: i32,
117}
118
119#[cfg(test)]
120mod tests {
121    use super::super::test_utils;
122    use super::*;
123
124    // Warning logs for initial LR and gamma are not tested because there seems no straightforward
125    // way to do it.
126    //
127    // Creating a mock logger that collects logs into `String` for later examination seems a possible
128    // solution, but unit tests run in the same process in parallel, where the single logger would
129    // be shared by multiple tests, so logs from different tests would be mixed up with no easy way
130    // to separate them.
131    // Using "--test-threads=1" could prevent mixup, but whether the ability to test logging is
132    // worth the slowdown would be a question. Also, using a primitive provided by `std` to
133    // synchronize the logger across tests is not an option since we need to support `no-std`.
134    // Maybe the mocking approach can be reconsidered after we are given an option to run tests in
135    // separate processes like what the issue below is proposing:
136    //     https://github.com/rust-lang/rust/issues/47506
137    //
138    // As a side note, a helper crate exists for the exact purpose:
139    //     https://crates.io/crates/testing_logger
140    // but the crate has been unmaintained and using it would introduce another dependency.
141
142    #[test]
143    fn test_config_step_size_zero() {
144        let r = StepLrSchedulerConfig::new(1.0, 0).build();
145        assert!(r.is_err(), "Should return an error");
146    }
147
148    #[test]
149    fn test_config_step_size_nonzero() {
150        let r = StepLrSchedulerConfig::new(1.0, 1).build();
151        assert!(r.is_ok(), "Should return a success value");
152    }
153
154    #[test]
155    fn test_config_default_gamma() {
156        const INIT_LR: LearningRate = 0.4;
157        const STEP_SIZE: usize = 2;
158
159        let mut default = StepLrSchedulerConfig::new(INIT_LR, STEP_SIZE)
160            .build()
161            .unwrap();
162        let mut explicit = StepLrSchedulerConfig::new(INIT_LR, STEP_SIZE)
163            .with_gamma(0.1)
164            .build()
165            .unwrap();
166        test_utils::compare_steps(&mut default, &mut explicit, 3 * STEP_SIZE);
167    }
168
169    #[test]
170    fn test_lr_decreasing() {
171        let scheduler = StepLrSchedulerConfig::new(0.5, 3)
172            .with_gamma(0.1)
173            .build()
174            .unwrap();
175        let expected_lrs = [0.5, 0.5, 0.5, 0.05, 0.05, 0.05, 0.005, 0.005, 0.005];
176        test_utils::check_lr_sequence(scheduler, expected_lrs);
177    }
178
179    #[test]
180    fn test_lr_increasing() {
181        let scheduler = StepLrSchedulerConfig::new(0.1, 2)
182            .with_gamma(2.0)
183            .build()
184            .unwrap();
185        let expected_lrs = [0.1, 0.1, 0.2, 0.2, 0.4, 0.4];
186        test_utils::check_lr_sequence(scheduler, expected_lrs);
187    }
188
189    #[test]
190    fn test_lr_unchanging() {
191        let scheduler = StepLrSchedulerConfig::new(3.1, 1)
192            .with_gamma(1.0)
193            .build()
194            .unwrap();
195        let expected_lrs = [3.1, 3.1, 3.1];
196        test_utils::check_lr_sequence(scheduler, expected_lrs);
197    }
198
199    #[test]
200    fn test_save_and_load() {
201        const STEP_SIZE: usize = 10;
202
203        let scheduler = StepLrSchedulerConfig::new(0.007, STEP_SIZE)
204            .with_gamma(0.03)
205            .build()
206            .unwrap();
207        test_utils::check_save_load(scheduler, 3 * STEP_SIZE / 2);
208    }
209
210    // It's too time consuming to actually run a scheduler `i32::MAX` steps, so an approach that
211    // depends on private fields is used to implement the test.
212    #[test]
213    fn test_number_of_calls_within_limit() {
214        // Create a scheduler that has already run `i32::MAX` steps
215        let mut scheduler = StepLrSchedulerConfig::new(0.1, 2).build().unwrap();
216        scheduler.load_record(LrSchedulerRecord::from_state(&StepLrSchedulerState {
217            iter_idx: i32::MAX - 1,
218        }));
219        scheduler.step();
220    }
221
222    #[test]
223    #[should_panic = "i32::MAX"]
224    fn test_number_of_calls_over_limit() {
225        // Create a scheduler that has already run `i32::MAX` steps
226        let mut scheduler = StepLrSchedulerConfig::new(0.1, 2).build().unwrap();
227        scheduler.load_record(LrSchedulerRecord::from_state(&StepLrSchedulerState {
228            iter_idx: i32::MAX - 1,
229        }));
230        scheduler.step();
231        scheduler.step();
232    }
233}
burn_optim/lr_scheduler/step.rs

burn_optim/lr_scheduler/
step.rs
