optirs_core/schedulers/
mod.rs

1// Learning rate schedulers for optimizers
2//
3// This module provides various learning rate schedulers that adjust the learning rate
4// of optimizers during training based on different strategies.
5
6use scirs2_core::ndarray::{Dimension, ScalarOperand};
7use scirs2_core::numeric::Float;
8use std::fmt::Debug;
9
10use crate::optimizers::Optimizer;
11
12/// Trait for learning rate schedulers
13pub trait LearningRateScheduler<A: Float + Debug + ScalarOperand> {
14    /// Get the learning rate at the current step
15    fn get_learning_rate(&self) -> A;
16
17    /// Update the scheduler state and return the new learning rate
18    fn step(&mut self) -> A;
19
20    /// Update the scheduler state using an observed metric and return the new learning rate
21    ///
22    /// The default implementation ignores `metric` and forwards to
23    /// [`LearningRateScheduler::step`], which is the correct behaviour for every schedule
24    /// that is driven purely by the step counter. Metric-driven schedulers such as
25    /// [`ReduceOnPlateau`] override this method with the real plateau logic.
26    ///
27    /// The method is deliberately non-generic so the trait stays object safe and can keep
28    /// being used as `Box<dyn LearningRateScheduler<A>>`.
29    fn step_with_metric(&mut self, metric: A) -> A {
30        let _ = metric;
31        self.step()
32    }
33
34    /// Apply the scheduler to an optimizer
35    fn apply_to<D: Dimension, O: Optimizer<A, D>>(&self, optimizer: &mut O)
36    where
37        Self: Sized,
38    {
39        optimizer.set_learning_rate(self.get_learning_rate());
40    }
41
42    /// Reset the scheduler state
43    fn reset(&mut self);
44}
45
46mod attention_aware;
47mod constant;
48mod cosine_annealing;
49mod cosine_annealing_warm_restarts;
50mod curriculum;
51mod custom_scheduler;
52mod cyclic_lr;
53mod exponential_decay;
54mod linear_decay;
55mod linear_warmup_decay;
56mod noise_injection;
57mod one_cycle;
58mod reduce_on_plateau;
59mod step_decay;
60mod vit_layer_decay;
61
62// Re-export schedulers
63pub use attention_aware::{AttentionAwareScheduler, TransformerComponentType};
64pub use constant::ConstantScheduler;
65pub use cosine_annealing::CosineAnnealing;
66pub use cosine_annealing_warm_restarts::CosineAnnealingWarmRestarts;
67pub use curriculum::{CurriculumScheduler, CurriculumStage, TransitionStrategy};
68pub use custom_scheduler::{CombinedScheduler, CustomScheduler, SchedulerBuilder};
69pub use cyclic_lr::{CyclicLR, CyclicMode};
70pub use exponential_decay::ExponentialDecay;
71pub use linear_decay::LinearDecay;
72pub use linear_warmup_decay::{DecayStrategy, LinearWarmupDecay};
73pub use noise_injection::{NoiseDistribution, NoiseInjectionScheduler};
74pub use one_cycle::{AnnealStrategy, OneCycle};
75pub use reduce_on_plateau::ReduceOnPlateau;
76pub use step_decay::StepDecay;
77pub use vit_layer_decay::ViTLayerDecay;
optirs_core/schedulers/mod.rs

optirs_core/schedulers/
mod.rs
