rust_lstm/
optimizers.rs

1use crate::schedulers::LearningRateScheduler;
2use ndarray::Array2;
3use std::collections::HashMap;
4
5/// Optimizer trait for parameter updates during training
6pub trait Optimizer {
7    fn update(&mut self, param_id: &str, param: &mut Array2<f64>, gradient: &Array2<f64>);
8    fn reset(&mut self);
9
10    /// Set the learning rate dynamically (for compatibility with schedulers)
11    fn set_learning_rate(&mut self, lr: f64);
12
13    /// Get the current learning rate
14    fn get_learning_rate(&self) -> f64;
15}
16
17/// Stochastic Gradient Descent optimizer
18pub struct SGD {
19    learning_rate: f64,
20}
21
22impl SGD {
23    pub fn new(learning_rate: f64) -> Self {
24        SGD { learning_rate }
25    }
26}
27
28impl Optimizer for SGD {
29    fn update(&mut self, _param_id: &str, param: &mut Array2<f64>, gradient: &Array2<f64>) {
30        *param = &*param - self.learning_rate * gradient;
31    }
32
33    fn reset(&mut self) {
34        // SGD has no state to reset
35    }
36
37    fn set_learning_rate(&mut self, lr: f64) {
38        self.learning_rate = lr;
39    }
40
41    fn get_learning_rate(&self) -> f64 {
42        self.learning_rate
43    }
44}
45
46/// Adam optimizer with adaptive learning rates (Kingma & Ba, 2015).
47///
48/// Moments and the bias-correction step count are kept per parameter, so each
49/// parameter follows the algorithm exactly however many parameters a model has.
50pub struct Adam {
51    learning_rate: f64,
52    beta1: f64,
53    beta2: f64,
54    epsilon: f64,
55    state: HashMap<String, AdamState>,
56}
57
58struct AdamState {
59    step: i32,
60    m: Array2<f64>,
61    v: Array2<f64>,
62}
63
64impl Adam {
65    pub fn new(learning_rate: f64) -> Self {
66        Adam::with_params(learning_rate, 0.9, 0.999, 1e-8)
67    }
68
69    pub fn with_params(learning_rate: f64, beta1: f64, beta2: f64, epsilon: f64) -> Self {
70        Adam {
71            learning_rate,
72            beta1,
73            beta2,
74            epsilon,
75            state: HashMap::new(),
76        }
77    }
78}
79
80impl Optimizer for Adam {
81    fn update(&mut self, param_id: &str, param: &mut Array2<f64>, gradient: &Array2<f64>) {
82        let state = self
83            .state
84            .entry(param_id.to_string())
85            .or_insert_with(|| AdamState {
86                step: 0,
87                m: Array2::zeros(param.raw_dim()),
88                v: Array2::zeros(param.raw_dim()),
89            });
90        state.step += 1;
91        state.m = self.beta1 * &state.m + (1.0 - self.beta1) * gradient;
92        state.v = self.beta2 * &state.v + (1.0 - self.beta2) * gradient * gradient;
93
94        let m_hat = &state.m / (1.0 - self.beta1.powi(state.step));
95        let v_hat = &state.v / (1.0 - self.beta2.powi(state.step));
96
97        let update = self.learning_rate * m_hat / (v_hat.mapv(f64::sqrt) + self.epsilon);
98        *param = &*param - update;
99    }
100
101    fn reset(&mut self) {
102        self.state.clear();
103    }
104
105    fn set_learning_rate(&mut self, lr: f64) {
106        self.learning_rate = lr;
107    }
108
109    fn get_learning_rate(&self) -> f64 {
110        self.learning_rate
111    }
112}
113
114/// RMSprop optimizer
115pub struct RMSprop {
116    learning_rate: f64,
117    alpha: f64,
118    epsilon: f64,
119    v: HashMap<String, Array2<f64>>,
120}
121
122impl RMSprop {
123    pub fn new(learning_rate: f64) -> Self {
124        RMSprop::with_params(learning_rate, 0.99, 1e-8)
125    }
126
127    pub fn with_params(learning_rate: f64, alpha: f64, epsilon: f64) -> Self {
128        RMSprop {
129            learning_rate,
130            alpha,
131            epsilon,
132            v: HashMap::new(),
133        }
134    }
135}
136
137impl Optimizer for RMSprop {
138    fn update(&mut self, param_id: &str, param: &mut Array2<f64>, gradient: &Array2<f64>) {
139        if !self.v.contains_key(param_id) {
140            self.v
141                .insert(param_id.to_string(), Array2::zeros(param.raw_dim()));
142        }
143
144        let v_t = self.v.get_mut(param_id).unwrap();
145
146        *v_t = self.alpha * &*v_t + (1.0 - self.alpha) * gradient * gradient;
147
148        let update = self.learning_rate * gradient / (v_t.map(|x| x.sqrt()) + self.epsilon);
149        *param = &*param - update;
150    }
151
152    fn reset(&mut self) {
153        self.v.clear();
154    }
155
156    fn set_learning_rate(&mut self, lr: f64) {
157        self.learning_rate = lr;
158    }
159
160    fn get_learning_rate(&self) -> f64 {
161        self.learning_rate
162    }
163}
164
165/// Wrapper that combines an optimizer with a learning rate scheduler
166pub struct ScheduledOptimizer<O: Optimizer, S: LearningRateScheduler> {
167    optimizer: O,
168    scheduler: S,
169    base_lr: f64,
170    current_epoch: usize,
171}
172
173impl<O: Optimizer, S: LearningRateScheduler> ScheduledOptimizer<O, S> {
174    pub fn new(optimizer: O, scheduler: S, base_lr: f64) -> Self {
175        ScheduledOptimizer {
176            optimizer,
177            scheduler,
178            base_lr,
179            current_epoch: 0,
180        }
181    }
182
183    /// Step the scheduler (should be called at the end of each epoch)
184    pub fn step(&mut self) {
185        self.current_epoch += 1;
186        let new_lr = self.scheduler.get_lr(self.current_epoch, self.base_lr);
187        self.optimizer.set_learning_rate(new_lr);
188    }
189
190    /// Step with validation loss (for ReduceLROnPlateau)
191    pub fn step_with_val_loss(&mut self, val_loss: f64) {
192        self.current_epoch += 1;
193        // For ReduceLROnPlateau, we need special handling
194        let base_lr = self.base_lr; // Copy the value before mutable borrow
195        let new_lr = if let Some(plateau_scheduler) = self.scheduler_as_plateau_mut() {
196            plateau_scheduler.step(val_loss, base_lr)
197        } else {
198            self.scheduler.get_lr(self.current_epoch, self.base_lr)
199        };
200        self.optimizer.set_learning_rate(new_lr);
201    }
202
203    /// Get the current learning rate
204    pub fn get_current_lr(&self) -> f64 {
205        self.optimizer.get_learning_rate()
206    }
207
208    /// Get the current epoch
209    pub fn get_current_epoch(&self) -> usize {
210        self.current_epoch
211    }
212
213    /// Reset both optimizer and scheduler
214    pub fn reset(&mut self) {
215        self.optimizer.reset();
216        self.scheduler.reset();
217        self.current_epoch = 0;
218        self.optimizer.set_learning_rate(self.base_lr);
219    }
220
221    /// Get the scheduler name for logging
222    pub fn scheduler_name(&self) -> &'static str {
223        self.scheduler.name()
224    }
225
226    /// Helper method to downcast scheduler to ReduceLROnPlateau if possible
227    fn scheduler_as_plateau_mut(&mut self) -> Option<&mut crate::schedulers::ReduceLROnPlateau> {
228        // This is a bit of a hack since we can't downcast traits easily in Rust
229        // In practice, users should use step_with_val_loss only with ReduceLROnPlateau
230        // For now, we'll return None and let the caller handle it properly
231        None
232    }
233}
234
235impl<O: Optimizer, S: LearningRateScheduler> Optimizer for ScheduledOptimizer<O, S> {
236    fn update(&mut self, param_id: &str, param: &mut Array2<f64>, gradient: &Array2<f64>) {
237        self.optimizer.update(param_id, param, gradient);
238    }
239
240    fn reset(&mut self) {
241        self.reset(); // Call our custom reset that handles both optimizer and scheduler
242    }
243
244    fn set_learning_rate(&mut self, lr: f64) {
245        self.base_lr = lr;
246        self.optimizer.set_learning_rate(lr);
247    }
248
249    fn get_learning_rate(&self) -> f64 {
250        self.optimizer.get_learning_rate()
251    }
252}
253
254/// Helper functions to create common optimizer-scheduler combinations
255impl<O: Optimizer> ScheduledOptimizer<O, crate::schedulers::ConstantLR> {
256    pub fn constant(optimizer: O, lr: f64) -> Self {
257        Self::new(optimizer, crate::schedulers::ConstantLR, lr)
258    }
259}
260
261impl<O: Optimizer> ScheduledOptimizer<O, crate::schedulers::StepLR> {
262    pub fn step_lr(optimizer: O, lr: f64, step_size: usize, gamma: f64) -> Self {
263        Self::new(
264            optimizer,
265            crate::schedulers::StepLR::new(step_size, gamma),
266            lr,
267        )
268    }
269}
270
271impl<O: Optimizer> ScheduledOptimizer<O, crate::schedulers::ExponentialLR> {
272    pub fn exponential(optimizer: O, lr: f64, gamma: f64) -> Self {
273        Self::new(optimizer, crate::schedulers::ExponentialLR::new(gamma), lr)
274    }
275}
276
277impl<O: Optimizer> ScheduledOptimizer<O, crate::schedulers::CosineAnnealingLR> {
278    pub fn cosine_annealing(optimizer: O, lr: f64, t_max: usize, eta_min: f64) -> Self {
279        Self::new(
280            optimizer,
281            crate::schedulers::CosineAnnealingLR::new(t_max, eta_min),
282            lr,
283        )
284    }
285}
286
287impl<O: Optimizer> ScheduledOptimizer<O, crate::schedulers::PolynomialLR> {
288    pub fn polynomial(optimizer: O, lr: f64, total_iters: usize, power: f64, end_lr: f64) -> Self {
289        Self::new(
290            optimizer,
291            crate::schedulers::PolynomialLR::new(total_iters, power, end_lr),
292            lr,
293        )
294    }
295}
296
297impl<O: Optimizer> ScheduledOptimizer<O, crate::schedulers::CyclicalLR> {
298    pub fn cyclical(optimizer: O, base_lr: f64, max_lr: f64, step_size: usize) -> Self {
299        Self::new(
300            optimizer,
301            crate::schedulers::CyclicalLR::new(base_lr, max_lr, step_size),
302            base_lr,
303        )
304    }
305
306    pub fn cyclical_triangular2(optimizer: O, base_lr: f64, max_lr: f64, step_size: usize) -> Self {
307        let scheduler = crate::schedulers::CyclicalLR::new(base_lr, max_lr, step_size)
308            .with_mode(crate::schedulers::CyclicalMode::Triangular2);
309        Self::new(optimizer, scheduler, base_lr)
310    }
311
312    pub fn cyclical_exp_range(
313        optimizer: O,
314        base_lr: f64,
315        max_lr: f64,
316        step_size: usize,
317        gamma: f64,
318    ) -> Self {
319        let scheduler = crate::schedulers::CyclicalLR::new(base_lr, max_lr, step_size)
320            .with_mode(crate::schedulers::CyclicalMode::ExpRange)
321            .with_gamma(gamma);
322        Self::new(optimizer, scheduler, base_lr)
323    }
324}
325
326impl<O: Optimizer> ScheduledOptimizer<O, crate::schedulers::OneCycleLR> {
327    pub fn one_cycle(optimizer: O, max_lr: f64, total_steps: usize) -> Self {
328        Self::new(
329            optimizer,
330            crate::schedulers::OneCycleLR::new(max_lr, total_steps),
331            max_lr,
332        )
333    }
334}
335
336#[cfg(test)]
337mod tests {
338    use super::*;
339    use ndarray::arr2;
340
341    #[test]
342    fn test_sgd_optimizer() {
343        let mut optimizer = SGD::new(0.1);
344        let mut param = arr2(&[[1.0, 2.0], [3.0, 4.0]]);
345        let gradient = arr2(&[[0.1, 0.2], [0.3, 0.4]]);
346
347        let original_param = param.clone();
348        optimizer.update("test_param", &mut param, &gradient);
349
350        let expected = &original_param - 0.1 * &gradient;
351        assert!((param - expected).map(|x| x.abs()).sum() < 1e-10);
352    }
353
354    #[test]
355    fn test_adam_optimizer() {
356        let mut optimizer = Adam::new(0.001);
357        let mut param = arr2(&[[1.0, 2.0], [3.0, 4.0]]);
358        let gradient = arr2(&[[0.1, 0.2], [0.3, 0.4]]);
359
360        let original_param = param.clone();
361        optimizer.update("test_param", &mut param, &gradient);
362
363        assert!((param - original_param).map(|x| x.abs()).sum() > 1e-10);
364    }
365
366    #[test]
367    fn adam_steps_each_parameter_independently() {
368        let lr = 0.001;
369        let gradients = [
370            arr2(&[[0.1, -0.2], [0.3, 0.05]]),
371            arr2(&[[-0.4, 0.1], [0.2, 0.2]]),
372            arr2(&[[0.3, 0.3], [-0.1, 0.6]]),
373        ];
374
375        let mut alone = Adam::new(lr);
376        let mut a = arr2(&[[1.0, 2.0], [3.0, 4.0]]);
377        for g in &gradients {
378            alone.update("a", &mut a, g);
379        }
380
381        let mut shared = Adam::new(lr);
382        let mut b = arr2(&[[1.0, 2.0], [3.0, 4.0]]);
383        let mut others = [Array2::zeros((3, 3)), Array2::zeros((1, 5))];
384        for g in &gradients {
385            for (i, other) in others.iter_mut().enumerate() {
386                let og = Array2::from_elem(other.raw_dim(), 0.5);
387                shared.update(&format!("other_{i}"), other, &og);
388            }
389            shared.update("b", &mut b, g);
390        }
391
392        assert!((&a - &b).iter().all(|d| d.abs() < 1e-15));
393    }
394
395    #[test]
396    fn adam_first_step_moves_each_weight_by_the_learning_rate() {
397        let mut optimizer = Adam::new(0.01);
398        let mut first = arr2(&[[0.0, 0.0]]);
399        let mut second = arr2(&[[0.0, 0.0]]);
400        optimizer.update("first", &mut first, &arr2(&[[2.0, -0.5]]));
401        optimizer.update("second", &mut second, &arr2(&[[0.3, -7.0]]));
402        for value in first.iter().chain(second.iter()) {
403            assert!((value.abs() - 0.01).abs() < 1e-8);
404        }
405    }
406
407    #[test]
408    fn test_rmsprop_optimizer() {
409        let mut optimizer = RMSprop::new(0.01);
410        let mut param = arr2(&[[1.0, 2.0], [3.0, 4.0]]);
411        let gradient = arr2(&[[0.1, 0.2], [0.3, 0.4]]);
412
413        let original_param = param.clone();
414        optimizer.update("test_param", &mut param, &gradient);
415
416        assert!((param - original_param).map(|x| x.abs()).sum() > 1e-10);
417    }
418}
rust_lstm/optimizers.rs

rust_lstm/
optimizers.rs
