burn_optim/lr_scheduler/
cosine.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 a [Cosine Annealing learning rate
10/// scheduler](CosineAnnealingLrScheduler).
11///
12/// This scheduler uses a cosine annealing schedule without warm restarts,
13/// where the learning rate follows a cosine curve from `initial_lr` down to
14/// `min_lr` over `num_iters` steps. After `num_iters` steps, the learning rate
15/// continues along the cosine curve without restarting.
16///
17/// This corresponds to PyTorch's `CosineAnnealingLR` and is based on the
18/// closed-form schedule proposed in [SGDR: Stochastic Gradient Descent with Warm
19/// Restarts](https://arxiv.org/abs/1608.03983).
20#[derive(Config, Debug)]
21pub struct CosineAnnealingLrSchedulerConfig {
22    // The initial learning rate.
23    initial_lr: LearningRate,
24    // The final learning rate.
25    #[config(default = 0.0)]
26    min_lr: LearningRate,
27    // The number of iterations for the learning rate to reach `min_lr` along the cosine curve.
28    num_iters: usize,
29}
30
31impl CosineAnnealingLrSchedulerConfig {
32    /// Initializes a [Cosine learning rate scheduler](CosineAnnealingLrScheduler).
33    pub(crate) fn build(&self) -> Result<CosineAnnealingLrScheduler, String> {
34        if self.initial_lr <= 0. || self.initial_lr > 1. {
35            return Err("Initial learning rate must be greater than 0 and at most 1".into());
36        }
37        if self.min_lr < 0.0 || self.min_lr > self.initial_lr {
38            return Err(
39                "Minimum learning rate must be at least 0 and at most equal to the initial \
40                 learning rate"
41                    .into(),
42            );
43        }
44        if self.num_iters == 0 {
45            return Err("Number of iterations must be at least 1".into());
46        }
47
48        Ok(CosineAnnealingLrScheduler {
49            min_lr: self.min_lr,
50            max_lr: self.initial_lr,
51            num_iters: self.num_iters,
52            current_iter: usize::MAX,
53        })
54    }
55
56    /// Initializes a [module learning rate scheduler](ModuleLrScheduler).
57    ///
58    /// # Errors
59    ///
60    /// An error will be returned if any of the following conditions is true:
61    ///
62    /// * `initial_lr` is out of range (0.0, 1.0]
63    /// * `min_lr` is out of range [0.0, `initial_lr`]
64    /// * `num_iters` is 0
65    pub fn init(&self) -> Result<ModuleLrScheduler, String> {
66        self.build().map(|s| s.into())
67    }
68}
69
70/// A Cosine Annealing learning rate scheduler without warm restarts.
71///
72/// The learning rate follows the closed-form schedule proposed in
73/// [SGDR: Stochastic Gradient Descent with Warm
74/// Restarts](https://arxiv.org/abs/1608.03983), but without the periodic
75/// restarts. The iteration counter increases monotonically, so the learning
76/// rate continues along the cosine curve past `num_iters` without resetting.
77///
78/// See [CosineAnnealingLrSchedulerConfig] for configuration options.
79#[derive(Clone, Copy, Debug)]
80pub struct CosineAnnealingLrScheduler {
81    min_lr: LearningRate,
82    max_lr: LearningRate,
83    num_iters: usize,
84    current_iter: usize,
85}
86
87impl LrScheduler for CosineAnnealingLrScheduler {
88    fn step(&mut self) -> LearningRate {
89        // Make current_iter overflow from usize::MAX to 0 to get the initial learning rate on the
90        // first call. We could've used i64 with an initial value -1, but keeping it in usize saves
91        // us from some type casting here.
92        self.current_iter = self.current_iter.wrapping_add(1);
93        self.min_lr
94            + 0.5
95                * (self.max_lr - self.min_lr)
96                * (1.0
97                    + (self.current_iter as f64 / self.num_iters as f64 * std::f64::consts::PI)
98                        .cos())
99    }
100
101    fn to_record(&self) -> LrSchedulerRecord {
102        LrSchedulerRecord::from_state(&CosineAnnealingLrSchedulerState {
103            current_iter: self.current_iter,
104        })
105    }
106
107    fn load_record(&mut self, record: LrSchedulerRecord) {
108        if let Some(state) = record.into_state::<CosineAnnealingLrSchedulerState>() {
109            self.current_iter = state.current_iter;
110        }
111    }
112}
113
114/// The serializable state of a [cosine annealing scheduler](CosineAnnealingLrScheduler).
115#[derive(RecordState, Clone, Debug)]
116pub struct CosineAnnealingLrSchedulerState {
117    current_iter: usize,
118}
119
120#[cfg(test)]
121mod tests {
122    use super::super::test_utils;
123    use super::*;
124
125    #[test]
126    fn config_initial_lr_too_low() {
127        let r = CosineAnnealingLrSchedulerConfig::new(0., 10).build();
128        assert!(r.is_err(), "Should return an error");
129        assert_eq!(
130            r.unwrap_err(),
131            "Initial learning rate must be greater than 0 and at most 1",
132            "Error messages should match",
133        );
134    }
135
136    #[test]
137    fn config_initial_lr_too_high() {
138        let r = CosineAnnealingLrSchedulerConfig::new(1.5, 10).build();
139        assert!(r.is_err(), "Should return an error");
140        assert_eq!(
141            r.unwrap_err(),
142            "Initial learning rate must be greater than 0 and at most 1",
143            "Error messages should match",
144        );
145    }
146
147    #[test]
148    fn config_min_lr_too_low() {
149        let r = CosineAnnealingLrSchedulerConfig::new(0.5, 10)
150            .with_min_lr(-0.1)
151            .build();
152        assert!(r.is_err(), "Should return an error");
153        assert_eq!(
154            r.unwrap_err(),
155            "Minimum learning rate must be at least 0 and at most equal to the initial learning \
156             rate",
157            "Error messages should match",
158        );
159    }
160
161    #[test]
162    fn config_min_lr_too_high() {
163        let r = CosineAnnealingLrSchedulerConfig::new(0.5, 10)
164            .with_min_lr(0.6)
165            .build();
166        assert!(r.is_err(), "Should return an error");
167        assert_eq!(
168            r.unwrap_err(),
169            "Minimum learning rate must be at least 0 and at most equal to the initial learning \
170             rate",
171            "Error messages should match",
172        );
173    }
174
175    #[test]
176    fn config_num_iters_too_low() {
177        let r = CosineAnnealingLrSchedulerConfig::new(0.5, 0).build();
178        assert!(r.is_err(), "Should return an error");
179        assert_eq!(
180            r.unwrap_err(),
181            "Number of iterations must be at least 1",
182            "Error messages should match",
183        );
184    }
185
186    #[test]
187    fn test_lr_change() {
188        const INITIAL_LR: LearningRate = 0.5;
189        const MIN_LR: LearningRate = 0.1;
190
191        let scheduler = CosineAnnealingLrSchedulerConfig::new(INITIAL_LR, 2)
192            .with_min_lr(MIN_LR)
193            .build()
194            .unwrap();
195        let expected_lrs = [
196            INITIAL_LR,                  // cos(0)
197            (INITIAL_LR + MIN_LR) * 0.5, // cos(PI/2)
198            MIN_LR,                      // cos(PI)
199            (INITIAL_LR + MIN_LR) * 0.5, // cos(3PI/2)
200            INITIAL_LR,                  // cos(2PI)
201        ];
202        test_utils::check_lr_sequence(scheduler, expected_lrs);
203    }
204
205    #[test]
206    fn test_save_and_load() {
207        const NUM_ITERS: usize = 9;
208        let scheduler = CosineAnnealingLrSchedulerConfig::new(1.0, NUM_ITERS)
209            .build()
210            .unwrap();
211        test_utils::check_save_load(scheduler, NUM_ITERS / 3 * 2);
212    }
213}
burn_optim/lr_scheduler/cosine.rs

burn_optim/lr_scheduler/
cosine.rs
