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
47pub struct Adam {
48    learning_rate: f64,
49    beta1: f64,
50    beta2: f64,
51    epsilon: f64,
52    t: i32,
53    m: HashMap<String, Array2<f64>>,
54    v: HashMap<String, Array2<f64>>,
55}
56
57impl Adam {
58    pub fn new(learning_rate: f64) -> Self {
59        Adam::with_params(learning_rate, 0.9, 0.999, 1e-8)
60    }
61
62    pub fn with_params(learning_rate: f64, beta1: f64, beta2: f64, epsilon: f64) -> Self {
63        Adam {
64            learning_rate,
65            beta1,
66            beta2,
67            epsilon,
68            t: 0,
69            m: HashMap::new(),
70            v: HashMap::new(),
71        }
72    }
73}
74
75impl Optimizer for Adam {
76    fn update(&mut self, param_id: &str, param: &mut Array2<f64>, gradient: &Array2<f64>) {
77        self.t += 1;
78
79        if !self.m.contains_key(param_id) {
80            self.m
81                .insert(param_id.to_string(), Array2::zeros(param.raw_dim()));
82            self.v
83                .insert(param_id.to_string(), Array2::zeros(param.raw_dim()));
84        }
85
86        let m_t = self.m.get_mut(param_id).unwrap();
87        let v_t = self.v.get_mut(param_id).unwrap();
88
89        *m_t = self.beta1 * &*m_t + (1.0 - self.beta1) * gradient;
90        *v_t = self.beta2 * &*v_t + (1.0 - self.beta2) * gradient * gradient;
91
92        let m_hat = &*m_t / (1.0 - self.beta1.powi(self.t));
93        let v_hat = &*v_t / (1.0 - self.beta2.powi(self.t));
94
95        let update = self.learning_rate * m_hat / (v_hat.map(|x| x.sqrt()) + self.epsilon);
96        *param = &*param - update;
97    }
98
99    fn reset(&mut self) {
100        self.t = 0;
101        self.m.clear();
102        self.v.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 test_rmsprop_optimizer() {
368        let mut optimizer = RMSprop::new(0.01);
369        let mut param = arr2(&[[1.0, 2.0], [3.0, 4.0]]);
370        let gradient = arr2(&[[0.1, 0.2], [0.3, 0.4]]);
371
372        let original_param = param.clone();
373        optimizer.update("test_param", &mut param, &gradient);
374
375        assert!((param - original_param).map(|x| x.abs()).sum() > 1e-10);
376    }
377}
rust_lstm/optimizers.rs

rust_lstm/
optimizers.rs
