zyx_optim/
adam.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/// # Adaptive momentum estimation optimizer
8#[derive(Module)]
9#[cfg_attr(feature = "py", pyo3::pyclass)]
10pub struct Adam {
11    /// learning rate (default: 1e-3)
12    pub learning_rate: f32,
13    /// coefficients used for computing running averages of gradient and its square (default: (0.9, 0.999))
14    pub betas: (f32, f32),
15    /// term added to the denominator to improve numerical stability (default: 1e-8)
16    pub eps: f32,
17    /// weight decay (L2 penalty) (default: 0)
18    pub weight_decay: f32,
19    /// whether to use the AMSGrad variant of this algorithm from the paper On the Convergence of Adam and Beyond (default: false)
20    pub amsgrad: bool,
21    /// m
22    pub m: Vec<Tensor>,
23    /// v
24    pub v: Vec<Tensor>,
25    /// vm
26    pub vm: Vec<Tensor>,
27    /// t
28    pub t: usize,
29}
30
31impl Default for Adam {
32    fn default() -> Self {
33        Self {
34            learning_rate: 0.001,
35            betas: (0.9, 0.999),
36            eps: 1e-8,
37            weight_decay: 0.0,
38            amsgrad: false,
39            m: Vec::new(),
40            v: Vec::new(),
41            vm: Vec::new(),
42            t: 0,
43        }
44    }
45}
46
47impl Adam {
48    /// Updates parameters with gradients.
49    /// Number of parameters must be the same as number of gradients.
50    /// Gradients can be None, those are simply skipped.
51    pub fn update<'a>(
52        &mut self,
53        parameters: impl IntoIterator<Item = &'a mut Tensor>,
54        gradients: impl IntoIterator<Item = Tensor>,
55    ) {
56        use zyx::Scalar;
57        //let params: Vec<&mut Tensor> = parameters.into_iter().collect();
58        //let grads: Vec<Option<Tensor>> = gradients.into_iter().collect();
59
60        /*assert_eq!(
61            params.len(),
62            grads.len(),
63            "Number of parameters != number of gradients."
64        );*/
65
66        self.t += 1;
67
68        for (i, (param, mut grad)) in parameters.into_iter().zip(gradients).enumerate() {
69            if self.weight_decay != 0.0 {
70                grad = grad + &*param * self.weight_decay;
71            }
72            if let Some(m) = self.m.get_mut(i) {
73                *m = &*m * self.betas.0 + &grad * (1.0 - self.betas.0);
74            } else {
75                self.m.push(&grad * (1.0 - self.betas.0));
76            }
77            if let Some(v) = self.v.get_mut(i) {
78                *v = &*v * self.betas.1 + &grad * &grad * (1.0 - self.betas.1);
79            } else {
80                self.v.push(&grad * &grad * (1.0 - self.betas.1));
81            }
82            let mh = &self.m[i] / (1.0 - self.betas.0.pow(self.t as f32));
83            let vh = &self.v[i] / (1.0 - self.betas.1.pow(self.t as f32));
84            if self.amsgrad {
85                if let Some(vm) = self.vm.get_mut(i) {
86                    *vm = vm.cmplt(&vh).unwrap().where_(vh, &*vm).unwrap();
87                } else {
88                    self.vm.push(vh);
89                }
90                *param = (&*param - self.learning_rate * mh / (self.vm[i].sqrt() + self.eps))
91                    .cast(param.dtype());
92            } else {
93                *param = (&*param - self.learning_rate * mh / (vh.sqrt() + self.eps))
94                    .cast(param.dtype());
95            }
96        }
97    }
98}
zyx_optim/adam.rs

zyx_optim/
adam.rs
