tensorlogic_train/optimizers/
adagrad.rs

1//! Adagrad optimizer (Adaptive Gradient).
2//!
3//! Adagrad adapts the learning rate for each parameter based on the historical
4//! sum of squared gradients, giving frequently occurring features lower learning rates.
5//!
6//! Reference: Duchi et al., "Adaptive Subgradient Methods for Online Learning
7//! and Stochastic Optimization", JMLR 2011
8
9use super::common::{compute_gradient_norm, GradClipMode, Optimizer, OptimizerConfig};
10use crate::{TrainError, TrainResult};
11use scirs2_core::ndarray::{Array, Ix2};
12use std::collections::HashMap;
13
14/// Adagrad optimizer (Adaptive Gradient).
15#[derive(Debug)]
16pub struct AdagradOptimizer {
17    config: OptimizerConfig,
18    /// Accumulated sum of squared gradients.
19    sum_squared_grads: HashMap<String, Array<f64, Ix2>>,
20}
21
22impl AdagradOptimizer {
23    /// Create a new Adagrad optimizer.
24    pub fn new(config: OptimizerConfig) -> Self {
25        Self {
26            config,
27            sum_squared_grads: HashMap::new(),
28        }
29    }
30
31    /// Apply gradient clipping if configured.
32    fn clip_gradients(&self, gradients: &mut HashMap<String, Array<f64, Ix2>>) {
33        if let Some(clip_value) = self.config.grad_clip {
34            match self.config.grad_clip_mode {
35                GradClipMode::Value => {
36                    for grad in gradients.values_mut() {
37                        grad.mapv_inplace(|g| g.max(-clip_value).min(clip_value));
38                    }
39                }
40                GradClipMode::Norm => {
41                    let total_norm = compute_gradient_norm(gradients);
42                    if total_norm > clip_value {
43                        let scale = clip_value / total_norm;
44                        for grad in gradients.values_mut() {
45                            grad.mapv_inplace(|g| g * scale);
46                        }
47                    }
48                }
49            }
50        }
51    }
52}
53
54impl Optimizer for AdagradOptimizer {
55    fn step(
56        &mut self,
57        parameters: &mut HashMap<String, Array<f64, Ix2>>,
58        gradients: &HashMap<String, Array<f64, Ix2>>,
59    ) -> TrainResult<()> {
60        let mut clipped_gradients = gradients.clone();
61        self.clip_gradients(&mut clipped_gradients);
62        let lr = self.config.learning_rate;
63        let eps = self.config.epsilon;
64        for (name, param) in parameters.iter_mut() {
65            let grad = clipped_gradients.get(name).ok_or_else(|| {
66                TrainError::OptimizerError(format!("Missing gradient for parameter: {}", name))
67            })?;
68            if !self.sum_squared_grads.contains_key(name) {
69                self.sum_squared_grads
70                    .insert(name.clone(), Array::zeros(param.raw_dim()));
71            }
72            let sum_sq = self
73                .sum_squared_grads
74                .get_mut(name)
75                .expect("sum_squared_grads initialized for all parameters");
76            let grad_squared = grad.mapv(|g| g * g);
77            *sum_sq = &*sum_sq + &grad_squared;
78            let update = grad / &sum_sq.mapv(|s| s.sqrt() + eps);
79            *param = &*param - &(update * lr);
80        }
81        Ok(())
82    }
83
84    fn zero_grad(&mut self) {}
85
86    fn get_lr(&self) -> f64 {
87        self.config.learning_rate
88    }
89
90    fn set_lr(&mut self, lr: f64) {
91        self.config.learning_rate = lr;
92    }
93
94    fn state_dict(&self) -> HashMap<String, Vec<f64>> {
95        let mut state = HashMap::new();
96        for (name, sum_sq) in &self.sum_squared_grads {
97            state.insert(
98                format!("sum_squared_grads_{}", name),
99                sum_sq.iter().copied().collect(),
100            );
101        }
102        state
103    }
104
105    fn load_state_dict(&mut self, state: HashMap<String, Vec<f64>>) {
106        for (key, values) in state {
107            if let Some(name) = key.strip_prefix("sum_squared_grads_") {
108                if let Some(sum_sq) = self.sum_squared_grads.get(name) {
109                    let shape = sum_sq.raw_dim();
110                    if let Ok(arr) = Array::from_shape_vec(shape, values) {
111                        self.sum_squared_grads.insert(name.to_string(), arr);
112                    }
113                }
114            }
115        }
116    }
117}
118
119#[cfg(test)]
120mod tests {
121    use super::*;
122    use scirs2_core::ndarray::array;
123
124    #[test]
125    fn test_adagrad_optimizer() {
126        let config = OptimizerConfig {
127            learning_rate: 0.1,
128            ..Default::default()
129        };
130        let mut optimizer = AdagradOptimizer::new(config);
131        let mut params = HashMap::new();
132        params.insert("w".to_string(), array![[1.0, 2.0]]);
133        let mut grads = HashMap::new();
134        grads.insert("w".to_string(), array![[0.1, 0.2]]);
135        optimizer.step(&mut params, &grads).expect("unwrap");
136        let w = params.get("w").expect("unwrap");
137        assert!(w[[0, 0]] < 1.0);
138        assert!(w[[0, 1]] < 2.0);
139    }
140}
tensorlogic_train/optimizers/adagrad.rs

tensorlogic_train/optimizers/
adagrad.rs
