zyx_optim/
sgd.rs

1// Copyright (C) 2025 zk4x
2// SPDX-License-Identifier: LGPL-3.0-only WITH Classpath-exception-2.0
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 = Tensor>,
49    ) {
50        let params: Vec<&mut Tensor> = parameters.into_iter().collect();
51        let grads: Vec<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, mut grad) in params.into_iter().zip(grads) {
61            if self.weight_decay != 0.0 {
62                grad = grad + param.clone() * self.weight_decay;
63            }
64            if self.momentum != 0.0 {
65                if bias_idx < self.bias.len() {
66                    self.bias[bias_idx] = self.bias[bias_idx].clone() * self.momentum
67                        + grad.clone() * (1.0 - self.dampening);
68                } else {
69                    self.bias.push(grad.clone());
70                }
71                if self.nesterov {
72                    grad = grad + self.bias[bias_idx].clone() * self.momentum;
73                } else {
74                    grad = self.bias[bias_idx].clone();
75                }
76                bias_idx += 1;
77            }
78            if self.maximize {
79                *param = (param.clone() + grad * self.learning_rate).cast(param.dtype());
80            } else {
81                *param = (param.clone() - grad * self.learning_rate).cast(param.dtype());
82            }
83        }
84    }
85}
zyx_optim/sgd.rs

zyx_optim/
sgd.rs
