zyx_optim/
sgd.rs

1// Copyright (C) 2025 zk4x
2// SPDX-License-Identifier: LGPL-3.0-only
3
4use zyx::Tensor;
5use zyx_derive::Module;
6
7/// # Stochastic gradient descent optimizer
8#[derive(Module)]
9#[cfg_attr(feature = "py", pyo3::pyclass)]
10pub struct SGD {
11    /// learning rate (default: 0.001)
12    pub learning_rate: f32,
13    /// momentum factor (default: 0.0)
14    pub momentum: f32,
15    /// weight decay (L2 penalty) (default: 0.0)
16    pub weight_decay: f32,
17    /// dampening for momentum (default: 0.0)
18    pub dampening: f32,
19    /// enables Nesterov momentum (default: false)
20    pub nesterov: bool,
21    /// maximize the objective with respect to the params, instead of minimizing (default: false)
22    pub maximize: bool,
23    /// stores momentum, starts empty and will be initialized on demand
24    pub bias: Vec<Tensor>,
25}
26
27impl Default for SGD {
28    fn default() -> Self {
29        Self {
30            learning_rate: 0.001,
31            momentum: 0.0,
32            weight_decay: 0.0,
33            dampening: 0.0,
34            nesterov: false,
35            maximize: false,
36            bias: Vec::new(),
37        }
38    }
39}
40
41impl SGD {
42    /// Updates parameters with gradients.
43    /// Number of parameters must be the same as number of gradients.
44    /// Gradients can be None, those are simply skipped.
45    pub fn update<'a>(
46        &mut self,
47        parameters: impl IntoIterator<Item = &'a mut Tensor>,
48        gradients: impl IntoIterator<Item = Option<Tensor>>,
49    ) {
50        let params: Vec<&mut Tensor> = parameters.into_iter().collect();
51        let grads: Vec<Option<Tensor>> = gradients.into_iter().collect();
52
53        assert_eq!(
54            params.len(),
55            grads.len(),
56            "Number of parameters != number of gradients."
57        );
58
59        let mut bias_idx = 0usize;
60        for (param, grad) in params.into_iter().zip(grads) {
61            if let Some(mut grad) = grad {
62                if self.weight_decay != 0.0 {
63                    grad = grad + param.clone() * self.weight_decay;
64                }
65                if self.momentum != 0.0 {
66                    if bias_idx < self.bias.len() {
67                        self.bias[bias_idx] =
68                            self.bias[bias_idx].clone() * self.momentum + grad.clone() * (1.0 - self.dampening);
69                    } else {
70                        self.bias.push(grad.clone());
71                    }
72                    if self.nesterov {
73                        grad = grad + self.bias[bias_idx].clone() * self.momentum;
74                    } else {
75                        grad = self.bias[bias_idx].clone();
76                    }
77                    bias_idx += 1;
78                }
79                if self.maximize {
80                    *param = (&*param + grad * self.learning_rate).cast(param.dtype());
81                } else {
82                    *param = (&*param - grad * self.learning_rate).cast(param.dtype());
83                }
84            }
85        }
86    }
87}
zyx_optim/sgd.rs

zyx_optim/
sgd.rs
