luma_optim/
sgd.rs

1use luma_io::lpk::LumaPack;
2use luma_tensor::{Device, GradStore, Scalar, Tensor, no_grad};
3
4use super::Optimizer;
5
6pub struct SGD<D: Device> {
7    pub params: Vec<Tensor<D>>,
8    pub learning_rate: f64,
9}
10
11impl<D: Device> SGD<D> {
12    pub fn new(params: impl Into<Vec<Tensor<D>>>, learning_rate: f64) -> Self {
13        Self { params: params.into(), learning_rate }
14    }
15}
16
17impl<D: Device> Optimizer for SGD<D> {
18    type Device = D;
19
20    fn get_lr(&self) -> f64 {
21        self.learning_rate
22    }
23
24    fn set_lr(&mut self, lr: f64) {
25        self.learning_rate = lr;
26    }
27
28    /// w_t = w_{t-1} - lr * g
29    fn step(&mut self, grads: &GradStore<Self::Device>) -> luma_tensor::Result<()> {
30        no_grad!();
31        for var in self.params.iter() {
32            if let Some(grad) = grads.get(var) {
33                var.sub_(&grad.mul_scalar(self.learning_rate)?)?;
34            }
35        }
36        Ok(())
37    }
38
39    fn state_dict(&self) -> luma_tensor::Result<LumaPack<Self::Device>> {
40        let mut pack = LumaPack::new();
41        pack.scalars.insert("lr".into(), Scalar::F64(self.learning_rate));
42        Ok(pack)
43    }
44
45    fn load_state_dict(&mut self, pack: &LumaPack<Self::Device>) -> luma_tensor::Result<()> {
46        if let Some(lr) = pack.scalars.get("lr").and_then(|s| s.to_f64()) {
47            self.learning_rate = lr;
48        }
49        Ok(())
50    }
51}
52
53pub struct SGDM<D: Device> {
54    pub params: Vec<SGDMParam<D>>,
55    pub learning_rate: f64,
56    pub momentum: f64,
57}
58
59pub struct SGDMParam<D: Device> {
60    pub param: Tensor<D>,
61    pub velocity: Tensor<D>,
62}
63
64impl<D: Device> Optimizer for SGDM<D> {
65    type Device = D;
66
67    fn get_lr(&self) -> f64 {
68        self.learning_rate
69    }
70
71    fn set_lr(&mut self, lr: f64) {
72        self.learning_rate = lr;
73    }
74
75    ///
76    /// v = m*v + grad
77    ///
78    /// w_t = w_{t-1} - lr * v
79    ///
80    fn step(&mut self, grads: &GradStore<Self::Device>) -> luma_tensor::Result<()> {
81        no_grad!();
82        for SGDMParam { param, velocity } in self.params.iter_mut() {
83            if let Some(grad) = grads.get(&param) {
84                // update v: v = m * v + grad
85                velocity.mul_scalar_(self.momentum)?;
86                velocity.add_(grad)?;
87                param.sub_(&velocity.mul_scalar(self.learning_rate)?)?;
88            }
89        }
90
91        Ok(())
92    }
93
94    fn state_dict(&self) -> luma_tensor::Result<LumaPack<Self::Device>> {
95        let mut pack = LumaPack::new();
96        for (i, p) in self.params.iter().enumerate() {
97            pack.tensors.insert(format!("{i}.velocity"), luma_tensor::DynTensor::Float(p.velocity.clone()));
98        }
99        pack.scalars.insert("lr".into(), Scalar::F64(self.learning_rate));
100        pack.scalars.insert("momentum".into(), Scalar::F64(self.momentum));
101        Ok(pack)
102    }
103
104    fn load_state_dict(&mut self, pack: &LumaPack<Self::Device>) -> luma_tensor::Result<()> {
105        if let Some(lr) = pack.scalars.get("lr").and_then(|s| s.to_f64()) {
106            self.learning_rate = lr;
107        }
108        if let Some(m) = pack.scalars.get("momentum").and_then(|s| s.to_f64()) {
109            self.momentum = m;
110        }
111        for (i, p) in self.params.iter_mut().enumerate() {
112            let key = format!("{i}.velocity");
113            if let Some(dt) = pack.tensors.get(&key) {
114                if let Some(src) = dt.as_float() {
115                    p.velocity.copy_(src)?;
116                }
117            }
118        }
119        Ok(())
120    }
121}
luma_optim/sgd.rs

luma_optim/
sgd.rs
