optirs_core/schedulers/
custom_scheduler.rs

1// Custom scheduler framework
2//
3// This module provides a flexible framework for creating custom learning rate schedulers
4// using closures and function combinators.
5
6use scirs2_core::ndarray::ScalarOperand;
7use scirs2_core::numeric::Float;
8use std::cell::RefCell;
9use std::fmt::Debug;
10use std::marker::PhantomData;
11use std::rc::Rc;
12
13use super::LearningRateScheduler;
14
15/// A custom scheduler that uses a closure to compute the learning rate
16pub struct CustomScheduler<A, F>
17where
18    A: Float + Debug + ScalarOperand,
19    F: FnMut(usize) -> A,
20{
21    /// Function to compute the learning rate, wrapped in RefCell for interior mutability
22    lr_func: Rc<RefCell<F>>,
23    /// Current step count
24    step_count: usize,
25    /// Phantom data for type parameter
26    _phantom: PhantomData<A>,
27}
28
29impl<A, F> CustomScheduler<A, F>
30where
31    A: Float + Debug + ScalarOperand,
32    F: FnMut(usize) -> A,
33{
34    /// Create a new custom scheduler with the given initial learning rate and computation function
35    ///
36    /// # Arguments
37    ///
38    /// * `initial_lr` - The initial learning rate (used only for documentation)
39    /// * `lr_func` - A function that takes the current step count and returns the learning rate
40    ///
41    /// # Example
42    ///
43    /// ```
44    /// use optirs_core::schedulers::{CustomScheduler, LearningRateScheduler};
45    ///
46    /// // Create a scheduler that reduces the learning rate by 10% every 10 steps
47    /// let mut scheduler = CustomScheduler::new(0.1, |step| {
48    ///     0.1 * 0.9_f64.powi((step / 10) as i32)
49    /// });
50    ///
51    /// assert_eq!(scheduler.get_learning_rate(), 0.1);
52    /// scheduler.step();
53    /// ```
54    pub fn new(_initial_lr: A, lrfunc: F) -> Self {
55        Self {
56            lr_func: Rc::new(RefCell::new(lrfunc)),
57            step_count: 0,
58            _phantom: PhantomData,
59        }
60    }
61
62    /// Get the current step count
63    pub fn get_step_count(&self) -> usize {
64        self.step_count
65    }
66}
67
68impl<A, F> LearningRateScheduler<A> for CustomScheduler<A, F>
69where
70    A: Float + Debug + ScalarOperand,
71    F: FnMut(usize) -> A,
72{
73    fn get_learning_rate(&self) -> A {
74        // Borrow the function mutably and call it
75        let mut func = self.lr_func.borrow_mut();
76        func(self.step_count)
77    }
78
79    fn step(&mut self) -> A {
80        self.step_count += 1;
81        self.get_learning_rate()
82    }
83
84    fn reset(&mut self) {
85        self.step_count = 0;
86    }
87}
88
89/// Scheduler combinator that allows combining multiple schedulers
90pub struct CombinedScheduler<A, F1, F2, C>
91where
92    A: Float + Debug + ScalarOperand,
93    F1: FnMut(usize) -> A,
94    F2: FnMut(usize) -> A,
95    C: FnMut(A, A) -> A,
96{
97    /// First scheduler
98    scheduler1: CustomScheduler<A, F1>,
99    /// Second scheduler
100    scheduler2: CustomScheduler<A, F2>,
101    /// Combinator function wrapped in RefCell for interior mutability
102    combinator: Rc<RefCell<C>>,
103}
104
105impl<A, F1, F2, C> CombinedScheduler<A, F1, F2, C>
106where
107    A: Float + Debug + ScalarOperand,
108    F1: FnMut(usize) -> A,
109    F2: FnMut(usize) -> A,
110    C: FnMut(A, A) -> A,
111{
112    /// Create a new combined scheduler
113    ///
114    /// # Arguments
115    ///
116    /// * `scheduler1` - The first scheduler
117    /// * `scheduler2` - The second scheduler
118    /// * `combinator` - A function that combines the learning rates from both schedulers
119    ///
120    /// # Example
121    ///
122    /// ```
123    /// use optirs_core::schedulers::{CustomScheduler, CombinedScheduler, LearningRateScheduler};
124    ///
125    /// // Create a scheduler that uses both exponential decay and cosine annealing
126    /// let exponential = CustomScheduler::new(0.1, |step| {
127    ///     0.1 * 0.9_f64.powi((step / 10) as i32)
128    /// });
129    ///
130    /// let cosine = CustomScheduler::new(0.1, |step| {
131    ///     let total_steps = 100;
132    ///     let min_lr = 0.001;
133    ///     let progress = std::f64::consts::PI * (step as f64) / (total_steps as f64);
134    ///     min_lr + (0.1 - min_lr) * (1.0 + progress.cos()) / 2.0
135    /// });
136    ///
137    /// let mut scheduler = CombinedScheduler::new(
138    ///     exponential,
139    ///     cosine,
140    ///     |lr1, lr2| (lr1 + lr2) / 2.0  // Average the learning rates
141    /// );
142    ///
143    /// assert_eq!(scheduler.get_learning_rate(), 0.1);
144    /// scheduler.step();
145    /// ```
146    pub fn new(
147        scheduler1: CustomScheduler<A, F1>,
148        scheduler2: CustomScheduler<A, F2>,
149        combinator: C,
150    ) -> Self {
151        Self {
152            scheduler1,
153            scheduler2,
154            combinator: Rc::new(RefCell::new(combinator)),
155        }
156    }
157}
158
159impl<A, F1, F2, C> LearningRateScheduler<A> for CombinedScheduler<A, F1, F2, C>
160where
161    A: Float + Debug + ScalarOperand,
162    F1: FnMut(usize) -> A,
163    F2: FnMut(usize) -> A,
164    C: FnMut(A, A) -> A,
165{
166    fn get_learning_rate(&self) -> A {
167        let lr1 = self.scheduler1.get_learning_rate();
168        let lr2 = self.scheduler2.get_learning_rate();
169
170        // Borrow the combinator function mutably and call it
171        let mut combinator = self.combinator.borrow_mut();
172        combinator(lr1, lr2)
173    }
174
175    fn step(&mut self) -> A {
176        self.scheduler1.step();
177        self.scheduler2.step();
178        self.get_learning_rate()
179    }
180
181    fn reset(&mut self) {
182        self.scheduler1.reset();
183        self.scheduler2.reset();
184    }
185}
186
187/// Builder for creating custom schedulers
188pub struct SchedulerBuilder<A>
189where
190    A: Float + Debug + ScalarOperand,
191{
192    initial_lr: A,
193}
194
195impl<A> SchedulerBuilder<A>
196where
197    A: Float + Debug + ScalarOperand,
198{
199    /// Create a new scheduler builder with the given initial learning rate
200    pub fn new(initiallr: A) -> Self {
201        Self {
202            initial_lr: initiallr,
203        }
204    }
205
206    /// Create a step decay scheduler
207    ///
208    /// # Arguments
209    ///
210    /// * `step_size` - The number of steps after which the learning rate is decayed
211    /// * `gamma` - The decay factor
212    pub fn step_decay(
213        self,
214        step_size: usize,
215        gamma: A,
216    ) -> CustomScheduler<A, impl FnMut(usize) -> A> {
217        let initial_lr = self.initial_lr;
218        CustomScheduler::new(initial_lr, move |step| {
219            let decay_factor = gamma.powi((step / step_size) as i32);
220            initial_lr * decay_factor
221        })
222    }
223
224    /// Create an exponential decay scheduler
225    ///
226    /// # Arguments
227    ///
228    /// * `gamma` - The decay factor
229    pub fn exponential_decay(self, gamma: A) -> CustomScheduler<A, impl FnMut(usize) -> A> {
230        let initial_lr = self.initial_lr;
231        CustomScheduler::new(initial_lr, move |step| initial_lr * gamma.powi(step as i32))
232    }
233
234    /// Create a linear decay scheduler
235    ///
236    /// # Arguments
237    ///
238    /// * `total_steps` - The total number of steps
239    /// * `final_lr` - The final learning rate
240    pub fn linear_decay(
241        self,
242        total_steps: usize,
243        final_lr: A,
244    ) -> CustomScheduler<A, impl FnMut(usize) -> A> {
245        let initial_lr = self.initial_lr;
246        let total_steps =
247            A::from(total_steps).expect("CustomScheduler: total_steps must fit in A (f32/f64)");
248        CustomScheduler::new(initial_lr, move |step| {
249            let step = A::from(step).expect("CustomScheduler: step must fit in A (f32/f64)");
250            if step >= total_steps {
251                final_lr
252            } else {
253                let progress = step / total_steps;
254                initial_lr + progress * (final_lr - initial_lr)
255            }
256        })
257    }
258
259    /// Create a cosine annealing scheduler
260    ///
261    /// # Arguments
262    ///
263    /// * `total_steps` - The total number of steps
264    /// * `min_lr` - The minimum learning rate
265    pub fn cosine_annealing(
266        self,
267        total_steps: usize,
268        min_lr: A,
269    ) -> CustomScheduler<A, impl FnMut(usize) -> A> {
270        let initial_lr = self.initial_lr;
271        let total_steps =
272            A::from(total_steps).expect("CustomScheduler: total_steps must fit in A (f32/f64)");
273        let pi = A::from(std::f64::consts::PI)
274            .expect("CustomScheduler: pi constant must fit in A (f32/f64)");
275        CustomScheduler::new(initial_lr, move |step| {
276            let step = A::from(step).expect("CustomScheduler: step must fit in A (f32/f64)");
277            if step >= total_steps {
278                min_lr
279            } else {
280                let progress = pi * step / total_steps;
281                min_lr + (initial_lr - min_lr) * (A::one() + progress.cos()) / (A::one() + A::one())
282            }
283        })
284    }
285
286    /// Create a cyclic learning rate scheduler
287    ///
288    /// # Arguments
289    ///
290    /// * `step_size` - The half cycle size
291    /// * `max_lr` - The maximum learning rate
292    /// * `mode` - The cycle mode (triangular, triangular2, or exp_range)
293    pub fn cyclic_lr(
294        self,
295        step_size: usize,
296        max_lr: A,
297        mode: CyclicMode<A>,
298    ) -> CustomScheduler<A, impl FnMut(usize) -> A> {
299        let min_lr = self.initial_lr;
300        let step_size =
301            A::from(step_size).expect("CustomScheduler: step_size must fit in A (f32/f64)");
302        let two = A::one() + A::one();
303
304        // Move mode into the closure
305        let mode_inner = mode;
306
307        CustomScheduler::new(min_lr, move |step| {
308            let step = A::from(step).expect("CustomScheduler: step must fit in A (f32/f64)");
309            let cycle = (step / (two * step_size)).floor();
310            let x = (step / step_size - two * cycle).abs();
311
312            let scale = match mode_inner {
313                CyclicMode::Triangular => A::one(),
314                CyclicMode::Triangular2 => A::one() / (two.powi(cycle.to_i32().unwrap_or(0))),
315                CyclicMode::ExpRange(gamma) => gamma.powi(step.to_i32().unwrap_or(0)),
316            };
317
318            min_lr + scale * (max_lr - min_lr) * (A::one() - x).max(A::zero())
319        })
320    }
321
322    /// Create a custom scheduler with a user-defined function
323    ///
324    /// # Arguments
325    ///
326    /// * `func` - A function that takes the current step count and returns the learning rate
327    pub fn custom<F>(self, func: F) -> CustomScheduler<A, F>
328    where
329        F: FnMut(usize) -> A,
330    {
331        CustomScheduler::new(self.initial_lr, func)
332    }
333}
334
335/// Cyclic learning rate modes
336#[derive(Debug, Clone, Copy)]
337pub enum CyclicMode<A: Float> {
338    /// Triangular mode
339    Triangular,
340    /// Triangular2 mode - Cycle amplitude is cut in half after each cycle
341    Triangular2,
342    /// ExpRange mode - Cycle amplitude is scaled by gamma^step
343    ExpRange(A),
344}
345
346#[cfg(test)]
347mod tests {
348    use super::*;
349    use approx::assert_relative_eq;
350
351    #[test]
352    fn test_custom_scheduler() {
353        let mut scheduler =
354            CustomScheduler::new(0.1f64, |step| 0.1 * 0.9f64.powi((step / 10) as i32));
355
356        assert_eq!(scheduler.get_learning_rate(), 0.1);
357        assert_eq!(scheduler.step(), 0.1);
358        assert_eq!(scheduler.step(), 0.1);
359
360        // After 10 steps, we should see decay
361        for _ in 0..8 {
362            scheduler.step();
363        }
364        assert_relative_eq!(scheduler.get_learning_rate(), 0.09, epsilon = 1e-10);
365    }
366
367    #[test]
368    fn test_combined_scheduler() {
369        let scheduler1 = CustomScheduler::new(0.1f64, |step| 0.1 * 0.9f64.powi((step / 10) as i32));
370
371        let scheduler2 = CustomScheduler::new(0.2f64, |step| 0.2 * 0.8f64.powi((step / 5) as i32));
372
373        let mut combined =
374            CombinedScheduler::new(scheduler1, scheduler2, |lr1, lr2| lr1 * 0.3 + lr2 * 0.7);
375
376        assert_relative_eq!(
377            combined.get_learning_rate(),
378            0.1 * 0.3 + 0.2 * 0.7,
379            epsilon = 1e-10
380        );
381        combined.step();
382        assert_relative_eq!(
383            combined.get_learning_rate(),
384            0.1 * 0.3 + 0.2 * 0.7,
385            epsilon = 1e-10
386        );
387
388        // After 5 steps, scheduler2 should decay
389        for _ in 0..4 {
390            combined.step();
391        }
392        assert_relative_eq!(
393            combined.get_learning_rate(),
394            0.1 * 0.3 + 0.2 * 0.8 * 0.7,
395            epsilon = 1e-10
396        );
397    }
398
399    #[test]
400    fn test_scheduler_builder() {
401        // Test step decay
402        let mut step_scheduler = SchedulerBuilder::new(0.1f64).step_decay(10, 0.5);
403        assert_eq!(step_scheduler.get_learning_rate(), 0.1);
404        for _ in 0..10 {
405            step_scheduler.step();
406        }
407        assert_relative_eq!(step_scheduler.get_learning_rate(), 0.05, epsilon = 1e-10);
408
409        // Test exponential decay
410        let mut exp_scheduler = SchedulerBuilder::new(0.1f64).exponential_decay(0.95);
411        assert_eq!(exp_scheduler.get_learning_rate(), 0.1);
412        exp_scheduler.step();
413        assert_relative_eq!(
414            exp_scheduler.get_learning_rate(),
415            0.1 * 0.95,
416            epsilon = 1e-10
417        );
418
419        // Test linear decay
420        let mut linear_scheduler = SchedulerBuilder::new(0.1f64).linear_decay(100, 0.01);
421        assert_eq!(linear_scheduler.get_learning_rate(), 0.1);
422        linear_scheduler.step();
423        assert_relative_eq!(
424            linear_scheduler.get_learning_rate(),
425            0.1 - 0.0009, // 0.1 + 1/100 * (0.01 - 0.1)
426            epsilon = 1e-10
427        );
428
429        // Test cosine annealing
430        let mut cosine_scheduler = SchedulerBuilder::new(0.1f64).cosine_annealing(100, 0.01);
431        assert_eq!(cosine_scheduler.get_learning_rate(), 0.1);
432        cosine_scheduler.step();
433        // Check that the first step is less than the initial
434        assert!(cosine_scheduler.get_learning_rate() < 0.1);
435        // And greater than the minimum
436        assert!(cosine_scheduler.get_learning_rate() > 0.01);
437    }
438}
optirs_core/schedulers/custom_scheduler.rs

optirs_core/schedulers/
custom_scheduler.rs
