burn_optim/lr_scheduler/
base.rs

1use std::fmt::Debug;
2
3use alloc::collections::BTreeMap;
4pub(super) use alloc::string::String;
5use alloc::vec::Vec;
6use burn_core as burn;
7use burn_core::config::Config;
8
9use crate::lr_scheduler::composed::ComposedLrSchedulerConfig;
10use crate::lr_scheduler::cosine::CosineAnnealingLrSchedulerConfig;
11use crate::lr_scheduler::exponential::ExponentialLrSchedulerConfig;
12use crate::lr_scheduler::linear::LinearLrSchedulerConfig;
13use crate::lr_scheduler::noam::NoamLrSchedulerConfig;
14use crate::lr_scheduler::sequential::SequentialLrSchedulerConfig;
15use crate::lr_scheduler::step::StepLrSchedulerConfig;
16use crate::{RecordState, StateSink, StateSource, join_path};
17use burn::store::RecordError;
18use burn::tensor::{Bytes, Device};
19use burn_pack::{Reader, Scalar, Writer};
20
21use crate::LearningRate;
22
23macro_rules! impl_from_for_scheduler {
24    ($($variant:ident($config:ident)),* $(,)?) => {
25        $(
26            impl From<$config> for LrSchedulerConfig {
27                fn from(config: $config) -> Self {
28                    LrSchedulerConfig::$variant(config)
29                }
30            }
31        )*
32    };
33}
34
35/// Learning rate scheduler defines how the learning rate will evolve during training.
36pub trait LrScheduler: LrSchedulerClone + Send + Sync {
37    /// Perform the scheduler step, potentially updating its state, and returning the effective
38    /// learning rate.
39    fn step(&mut self) -> LearningRate;
40
41    /// Get the current state of the scheduler as a [record](LrSchedulerRecord).
42    fn to_record(&self) -> LrSchedulerRecord;
43
44    /// Load the state of the scheduler from a [record](LrSchedulerRecord).
45    fn load_record(&mut self, record: LrSchedulerRecord);
46}
47
48/// Implements the clone of a boxed [`LrScheduler`].
49pub trait LrSchedulerClone {
50    /// Clones a boxed [`LrScheduler`].
51    fn clone_box(&self) -> Box<dyn LrScheduler>;
52}
53
54impl<T> LrSchedulerClone for T
55where
56    T: 'static + LrScheduler + Clone,
57{
58    fn clone_box(&self) -> Box<dyn LrScheduler> {
59        Box::new(self.clone())
60    }
61}
62
63impl Clone for Box<dyn LrScheduler> {
64    fn clone(&self) -> Box<dyn LrScheduler> {
65        self.as_ref().clone_box()
66    }
67}
68
69/// A wrapper over a dynamic [`LrScheduler`].
70#[derive(Clone)]
71pub struct DynLrScheduler {
72    scheduler: Box<dyn LrScheduler>,
73}
74
75impl DynLrScheduler {
76    /// Perform the scheduler step, potentially updating its state, and returning the effective
77    /// learning rate.
78    pub fn step(&mut self) -> LearningRate {
79        self.scheduler.step()
80    }
81
82    /// Get the current state of the scheduler as a [record](LrSchedulerRecord).
83    pub fn to_record(&self) -> LrSchedulerRecord {
84        self.scheduler.to_record()
85    }
86
87    /// Load the state of the scheduler from a [record](LrSchedulerRecord).
88    pub fn load_record(mut self, record: LrSchedulerRecord) -> Self {
89        self.scheduler.load_record(record);
90        self
91    }
92}
93
94impl<S> From<S> for DynLrScheduler
95where
96    S: LrScheduler + 'static,
97{
98    fn from(scheduler: S) -> Self {
99        Self {
100            scheduler: Box::new(scheduler),
101        }
102    }
103}
104
105/// The serialized state of a [learning rate scheduler](LrScheduler), stored in the
106/// [burnpack](burn_pack) format.
107///
108/// Scheduler state is just a handful of scalars (step counters, current learning rate), so the
109/// record holds named typed scalars and no tensors. Composed schedulers nest their children's
110/// records under an index prefix via [`with_record`](Self::with_record) / [`record`](Self::record).
111#[derive(Default, Clone, Debug)]
112pub struct LrSchedulerRecord {
113    scalars: BTreeMap<String, Scalar>,
114}
115
116impl LrSchedulerRecord {
117    /// Create an empty record.
118    pub fn new() -> Self {
119        Self::default()
120    }
121
122    /// Whether the record holds no scalars.
123    pub fn is_empty(&self) -> bool {
124        self.scalars.is_empty()
125    }
126
127    /// Store a scalar under `key`.
128    pub fn with_scalar<V: Into<Scalar>>(mut self, key: &str, value: V) -> Self {
129        self.scalars.insert(String::from(key), value.into());
130        self
131    }
132
133    /// Read the scalar stored under `key`, if present and of a compatible type.
134    pub fn scalar<V: TryFrom<Scalar>>(&self, key: &str) -> Option<V> {
135        self.scalars
136            .get(key)
137            .copied()
138            .and_then(|scalar| V::try_from(scalar).ok())
139    }
140
141    /// Merge a child `record`'s scalars under `prefix` (used to compose schedulers).
142    pub fn with_record(mut self, prefix: &str, record: LrSchedulerRecord) -> Self {
143        for (key, value) in record.scalars {
144            self.scalars.insert(join_path(prefix, &key), value);
145        }
146        self
147    }
148
149    /// Extract the child record previously merged under `prefix`.
150    pub fn record(&self, prefix: &str) -> LrSchedulerRecord {
151        let head = join_path(prefix, "");
152        let scalars = self
153            .scalars
154            .iter()
155            .filter_map(|(key, value)| {
156                key.strip_prefix(&head)
157                    .map(|stripped| (String::from(stripped), *value))
158            })
159            .collect();
160        LrSchedulerRecord { scalars }
161    }
162
163    /// Build a record from a scheduler [state](RecordState).
164    ///
165    /// Scheduler state is scalar-only, so this reuses the same [`RecordState`] decomposition as
166    /// optimizer states (it panics in debug builds if a tensor leaf is produced).
167    pub fn from_state<S: RecordState>(state: &S) -> Self {
168        let mut sink = StateSink::default();
169        state.state_flatten("", &mut sink);
170        debug_assert!(
171            sink.tensors.is_empty(),
172            "learning rate scheduler state is expected to be scalar-only"
173        );
174        Self {
175            scalars: sink.scalars.into_iter().collect(),
176        }
177    }
178
179    /// Reconstruct a scheduler [state](RecordState) from this record.
180    ///
181    /// Uses the default device; scheduler state is scalar-only so no tensor is ever allocated.
182    pub fn into_state<S: RecordState>(&self) -> Option<S> {
183        let mut source = StateSource::new(self.scalars.clone());
184        S::state_unflatten("", &mut source, &Device::default())
185    }
186
187    /// Serialize the record to an in-memory burnpack byte buffer.
188    pub fn into_bytes(self) -> Result<Bytes, RecordError> {
189        Ok(self.into_writer().into_bytes()?)
190    }
191
192    /// Reconstruct a record from an in-memory burnpack byte buffer.
193    pub fn from_bytes(bytes: Bytes) -> Result<Self, RecordError> {
194        let reader = Reader::from_bytes(bytes)?;
195        Ok(Self {
196            scalars: reader.scalars().clone(),
197        })
198    }
199
200    /// Save the record to a burnpack file on disk.
201    #[cfg(feature = "std")]
202    pub fn save<P: AsRef<std::path::Path>>(self, path: P) -> Result<(), RecordError> {
203        self.into_writer().write_to_file(path)?;
204        Ok(())
205    }
206
207    /// Load the record from a burnpack file on disk.
208    #[cfg(feature = "std")]
209    pub fn load<P: AsRef<std::path::Path>>(path: P) -> Result<Self, RecordError> {
210        let reader = Reader::from_file(path)?;
211        Ok(Self {
212            scalars: reader.scalars().clone(),
213        })
214    }
215
216    fn into_writer(self) -> Writer {
217        let mut writer = Writer::new(Vec::new());
218        for (key, value) in &self.scalars {
219            writer = writer.with_scalar(key, *value);
220        }
221        writer
222    }
223}
224
225/// An enum for possible learning rate scheduler configs.
226#[derive(Config, Debug)]
227pub enum LrSchedulerConfig {
228    /// A constant learning rate.
229    Constant(LearningRate),
230    /// A [`LinearLrSchedulerConfig`]
231    Linear(LinearLrSchedulerConfig),
232    /// A [`CosineAnnealingLrSchedulerConfig`]
233    Cosine(CosineAnnealingLrSchedulerConfig),
234    /// A [`ExponentialLrSchedulerConfig`]
235    Exponential(ExponentialLrSchedulerConfig),
236    /// A [`NoamLrSchedulerConfig`]
237    Noam(NoamLrSchedulerConfig),
238    /// A [`StepLrSchedulerConfig`]
239    Step(StepLrSchedulerConfig),
240    /// A [`ComposedLrSchedulerConfig`]
241    Composed(ComposedLrSchedulerConfig),
242    /// A [`SequentialLrSchedulerConfig`]
243    Sequential(SequentialLrSchedulerConfig),
244}
245
246impl LrSchedulerConfig {
247    pub(crate) fn build(&self) -> Result<DynLrScheduler, String> {
248        Ok(match self {
249            Self::Constant(lr) => (*lr).into(),
250            Self::Linear(config) => config.build()?.into(),
251            Self::Cosine(config) => config.build()?.into(),
252            Self::Exponential(config) => config.build()?.into(),
253            Self::Noam(config) => config.build()?.into(),
254            Self::Step(config) => config.build()?.into(),
255            Self::Composed(config) => config.build()?.into(),
256            Self::Sequential(config) => config.build()?.into(),
257        })
258    }
259}
260
261impl_from_for_scheduler!(
262    Constant(LearningRate),
263    Linear(LinearLrSchedulerConfig),
264    Cosine(CosineAnnealingLrSchedulerConfig),
265    Exponential(ExponentialLrSchedulerConfig),
266    Noam(NoamLrSchedulerConfig),
267    Step(StepLrSchedulerConfig),
268    Composed(ComposedLrSchedulerConfig),
269    Sequential(SequentialLrSchedulerConfig),
270);
271
272#[cfg(test)]
273pub(super) mod test_utils {
274    use super::*;
275
276    // A small tolerance for learning rate comparisons. Depending on how learning rates are
277    // computed, floating-point arithmetic error might exceed f64::EPSILON, so a larger value is
278    // used here.
279    const LOOSE_EPSILON: LearningRate = 1e-10;
280
281    pub fn check_lr_sequence<I, S>(mut scheduler: S, expected_lrs: I)
282    where
283        I: IntoIterator<Item = LearningRate>,
284        S: LrScheduler,
285    {
286        expected_lrs
287            .into_iter()
288            .enumerate()
289            .for_each(|(i, expected)| {
290                let lr = scheduler.step();
291                assert!(
292                    (lr - expected).abs() < LOOSE_EPSILON,
293                    "Scheduled learning rate {lr} is not approximately equal to the expected value \
294                     {expected} at step {i}",
295                );
296            });
297    }
298
299    // save_at_step is the number of steps to run the scheduler before saving and loading back its
300    // state.
301    pub fn check_save_load<S>(mut scheduler: S, save_at_step: usize)
302    where
303        S: Clone + LrScheduler,
304    {
305        let mut truth = scheduler.clone();
306        // Consume some steps before saving and loading back
307        (0..save_at_step).for_each(|_| {
308            truth.step();
309            scheduler.step();
310        });
311        let rec = scheduler.to_record();
312        scheduler.load_record(rec);
313
314        // Validate that the scheduler resumes from where it left off.
315        compare_steps(&mut scheduler, &mut truth, save_at_step);
316    }
317
318    // Check if two schedulers produce the same learning rate sequences over the specified number of
319    // steps.
320    pub fn compare_steps<S: LrScheduler>(a: &mut S, b: &mut S, num_steps: usize) {
321        (0..num_steps).for_each(|i| {
322            let lr_a = a.step();
323            let lr_b = b.step();
324            assert!(
325                (lr_a - lr_b).abs() < LOOSE_EPSILON,
326                "The two learning rates ({lr_a}, {lr_b}) at position {i} in the remaining \
327                 sequences are not approximately equal",
328            );
329        });
330    }
331}
burn_optim/lr_scheduler/base.rs

burn_optim/lr_scheduler/
base.rs
