luma_optim/
momentum.rs

1use luma_io::lpk::LumaPack;
2use luma_tensor::{Device, DynTensor, Scalar, Tensor, no_grad};
3
4use super::Optimizer;
5
6#[derive(Clone, Debug)]
7pub struct MomentumConfig {
8    pub lr: f64,
9    pub momentum: f64,
10    pub weight_decay: f64,
11    pub dampening: f64,
12    pub nesterov: bool,
13}
14
15impl Default for MomentumConfig {
16    fn default() -> Self {
17        Self { lr: 1e-3, momentum: 0.9, weight_decay: 0.0, dampening: 0.0, nesterov: false }
18    }
19}
20
21struct MomentumParam<D: Device> {
22    param: Tensor<D>,
23    velocity: Tensor<D>,
24}
25
26pub struct Momentum<D: Device> {
27    params: Vec<MomentumParam<D>>,
28    config: MomentumConfig,
29}
30
31impl<D: Device> Momentum<D> {
32    pub fn new(params: impl Into<Vec<Tensor<D>>>, config: MomentumConfig) -> luma_tensor::Result<Self> {
33        let params = params
34            .into()
35            .into_iter()
36            .map(|param| {
37                let velocity = param.zeros_like()?;
38                Ok(MomentumParam { param, velocity })
39            })
40            .collect::<luma_tensor::Result<Vec<_>>>()?;
41        Ok(Self { params, config })
42    }
43}
44
45impl<D: Device> Optimizer for Momentum<D> {
46    type Device = D;
47
48    fn get_lr(&self) -> f64 {
49        self.config.lr
50    }
51
52    fn set_lr(&mut self, lr: f64) {
53        self.config.lr = lr;
54    }
55
56    fn step(&mut self, grads: &luma_tensor::GradStore<Self::Device>) -> luma_tensor::Result<()> {
57        no_grad!();
58
59        let lr = self.config.lr;
60        let momentum = self.config.momentum;
61        let weight_decay = self.config.weight_decay;
62        let dampening = self.config.dampening;
63        let nesterov = self.config.nesterov;
64
65        for MomentumParam { param, velocity } in self.params.iter_mut() {
66            if let Some(g) = grads.get(&param) {
67                let mut g = g.clone();
68                if weight_decay != 0.0 {
69                    /*
70                        origin: w_t = w_{t-1} - lr*g
71                        weight_decay: w_t = (1-lr*eta)w_{t-1} - lr*g
72                        equal to: w_t = w_{t-1} - lr*(g + eta*w)
73                    */
74                    g.add_(&param.mul_scalar(weight_decay)?)?;
75                }
76
77                // v_new = momentum * v_old + (1 - dampening) * g
78                if momentum != 0.0 {
79                    if dampening != 0.0 {
80                        g.mul_scalar_(1. - dampening)?;
81                    }
82
83                    // v = m*v + g
84                    velocity.mul_scalar_(momentum)?;
85                    velocity.add_(&g)?;
86
87                    if nesterov {
88                        g.add_(&velocity.mul_scalar(momentum)?)?;
89                    } else {
90                        g = velocity.clone();
91                    }
92                }
93
94                param.sub_(&g.mul_scalar(lr)?)?;
95            }
96        }
97
98        Ok(())
99    }
100
101    fn state_dict(&self) -> luma_tensor::Result<LumaPack<Self::Device>> {
102        let mut pack = LumaPack::new();
103        for (i, p) in self.params.iter().enumerate() {
104            pack.tensors.insert(format!("{i}.velocity"), DynTensor::Float(p.velocity.clone()));
105        }
106        pack.scalars.insert("lr".into(), Scalar::F64(self.config.lr));
107        pack.scalars.insert("momentum".into(), Scalar::F64(self.config.momentum));
108        pack.scalars.insert("weight_decay".into(), Scalar::F64(self.config.weight_decay));
109        pack.scalars.insert("dampening".into(), Scalar::F64(self.config.dampening));
110        pack.scalars.insert("nesterov".into(), Scalar::Bool(self.config.nesterov));
111        Ok(pack)
112    }
113
114    fn load_state_dict(&mut self, pack: &LumaPack<Self::Device>) -> luma_tensor::Result<()> {
115        if let Some(v) = pack.scalars.get("lr").and_then(|s| s.to_f64()) {
116            self.config.lr = v;
117        }
118        if let Some(v) = pack.scalars.get("momentum").and_then(|s| s.to_f64()) {
119            self.config.momentum = v;
120        }
121        if let Some(v) = pack.scalars.get("weight_decay").and_then(|s| s.to_f64()) {
122            self.config.weight_decay = v;
123        }
124        if let Some(v) = pack.scalars.get("dampening").and_then(|s| s.to_f64()) {
125            self.config.dampening = v;
126        }
127        if let Some(v) = pack.scalars.get("nesterov").and_then(|s| s.to_bool()) {
128            self.config.nesterov = v;
129        }
130        for (i, p) in self.params.iter_mut().enumerate() {
131            let key = format!("{i}.velocity");
132            if let Some(dt) = pack.tensors.get(&key) {
133                if let Some(src) = dt.as_float() {
134                    p.velocity.copy_(src)?;
135                }
136            }
137        }
138        Ok(())
139    }
140}
luma_optim/momentum.rs

luma_optim/
momentum.rs
