luma_optim/
rms_prop.rs

1use luma_io::lpk::LumaPack;
2use luma_tensor::{Device, DynTensor, GradStore, Scalar, Tensor, no_grad};
3
4use super::Optimizer;
5
6#[derive(Clone, Debug)]
7pub struct RMSPropConfig {
8    pub lr: f64,
9    pub alpha: f64, // 平滑常数（通常设为 0.99）
10    pub eps: f64,   // 极小值（通常设为 1e-8）
11    pub weight_decay: f64,
12}
13
14impl Default for RMSPropConfig {
15    fn default() -> Self {
16        Self { lr: 1e-2, alpha: 0.99, eps: 1e-8, weight_decay: 0.0 }
17    }
18}
19
20struct RMSPropParam<D: Device> {
21    param: Tensor<D>,
22    square_avg: Tensor<D>,
23}
24
25pub struct RMSProp<D: Device> {
26    params: Vec<RMSPropParam<D>>,
27    config: RMSPropConfig,
28}
29
30impl<D: Device> RMSProp<D> {
31    pub fn new(params: impl Into<Vec<Tensor<D>>>, config: RMSPropConfig) -> luma_tensor::Result<Self> {
32        let params = params
33            .into()
34            .into_iter()
35            .map(|param| {
36                let square_avg = param.zeros_like()?;
37                Ok(RMSPropParam { param, square_avg })
38            })
39            .collect::<luma_tensor::Result<Vec<_>>>()?;
40        Ok(Self { params, config })
41    }
42}
43
44impl<D: Device> Optimizer for RMSProp<D> {
45    type Device = D;
46
47    fn set_lr(&mut self, lr: f64) {
48        self.config.lr = lr;
49    }
50
51    fn get_lr(&self) -> f64 {
52        self.config.lr
53    }
54
55    fn step(&mut self, grads: &GradStore<Self::Device>) -> luma_tensor::Result<()> {
56        no_grad!();
57        let lr = self.config.lr;
58        let alpha = self.config.alpha;
59        let eps = self.config.eps;
60        let weight_decay = self.config.weight_decay;
61
62        for RMSPropParam { param, square_avg } in self.params.iter_mut() {
63            if let Some(g) = grads.get(&param) {
64                let g = g.clone();
65
66                /*
67                    1. weight decay
68
69                    w_t = w_{t-1} - lr*g
70                    w_t = (1-lr*eta)w_{t-1} - lr*g
71                    w_t = w_{t-1} - lr*(g+eta*w_{t+1})
72
73                    g = g+eta*w_{t+1}
74                */
75                if weight_decay != 0.0 {
76                    g.add_(&param.mul_scalar(weight_decay)?)?;
77                }
78
79                // 2. update square_avg
80                // square_avg = alpha*square_avg + (1-alpha)*g^2
81                square_avg.mul_scalar_(alpha)?;
82                square_avg.add_(&g.pow(2.0)?.mul_scalar(1. - alpha)?)?;
83
84                // 3. update param
85                let denom = square_avg.sqrt()?.add_scalar(eps)?;
86                param.sub_(&g.div(&denom)?.mul_scalar(lr)?)?
87            }
88        }
89
90        Ok(())
91    }
92
93    fn state_dict(&self) -> luma_tensor::Result<LumaPack<Self::Device>> {
94        let mut pack = LumaPack::new();
95        for (i, p) in self.params.iter().enumerate() {
96            pack.tensors.insert(format!("{i}.square_avg"), DynTensor::Float(p.square_avg.clone()));
97        }
98        pack.scalars.insert("lr".into(), Scalar::F64(self.config.lr));
99        pack.scalars.insert("alpha".into(), Scalar::F64(self.config.alpha));
100        pack.scalars.insert("eps".into(), Scalar::F64(self.config.eps));
101        pack.scalars.insert("weight_decay".into(), Scalar::F64(self.config.weight_decay));
102        Ok(pack)
103    }
104
105    fn load_state_dict(&mut self, pack: &LumaPack<Self::Device>) -> luma_tensor::Result<()> {
106        if let Some(v) = pack.scalars.get("lr").and_then(|s| s.to_f64()) {
107            self.config.lr = v;
108        }
109        if let Some(v) = pack.scalars.get("alpha").and_then(|s| s.to_f64()) {
110            self.config.alpha = v;
111        }
112        if let Some(v) = pack.scalars.get("eps").and_then(|s| s.to_f64()) {
113            self.config.eps = v;
114        }
115        if let Some(v) = pack.scalars.get("weight_decay").and_then(|s| s.to_f64()) {
116            self.config.weight_decay = v;
117        }
118        for (i, p) in self.params.iter_mut().enumerate() {
119            let key = format!("{i}.square_avg");
120            if let Some(dt) = pack.tensors.get(&key) {
121                if let Some(src) = dt.as_float() {
122                    p.square_avg.copy_(src)?;
123                }
124            }
125        }
126        Ok(())
127    }
128}
luma_optim/rms_prop.rs

luma_optim/
rms_prop.rs
