burn_optim/optim/
rmsprop.rs

1use burn_core as burn;
2
3use crate::RecordState;
4
5use super::{
6    Optimizer,
7    decay::{WeightDecay, WeightDecayConfig},
8    module_optimizer::ModuleOptimizer,
9};
10use crate::{LearningRate, grad_clipping::GradientClippingConfig};
11
12use burn::config::Config;
13use burn::tensor::{Device, Tensor};
14
15/// Configuration to create the [RmsProp](RmsProp) optimizer.
16#[derive(Config, Debug)]
17pub struct RmsPropConfig {
18    /// Smoothing constant.
19    #[config(default = 0.99)]
20    alpha: f32,
21    /// momentum for RmsProp.
22    #[config(default = 0.9)]
23    momentum: f32,
24    /// A value required for numerical stability.
25    #[config(default = 1e-5)]
26    epsilon: f32,
27    /// if True, compute the centered RmsProp, the gradient is normalized by an estimation of its variance
28    #[config(default = false)]
29    centered: bool,
30    /// [Weight decay](WeightDecayConfig) config.
31    weight_decay: Option<WeightDecayConfig>,
32    /// [Gradient Clipping](GradientClippingConfig) config.
33    grad_clipping: Option<GradientClippingConfig>,
34}
35
36impl RmsPropConfig {
37    /// Build a [`RmsProp`] from the config.
38    ///
39    /// The bare optimizer, which
40    /// [`ModuleOptimizer::with_group`](crate::ModuleOptimizer::with_group) takes to
41    /// optimize one parameter group. [`init`](Self::init) is the whole-module
42    /// counterpart, and the only one that applies the configured gradient clipping.
43    pub fn build(&self) -> RmsProp {
44        let weight_decay = self.weight_decay.as_ref().map(WeightDecay::new);
45        RmsProp {
46            alpha: self.alpha,
47            centered: self.centered,
48            weight_decay,
49            momentum: RmsPropMomentum {
50                momentum: self.momentum,
51                epsilon: self.epsilon,
52            },
53        }
54    }
55
56    /// Initialize RmsProp optimizer.
57    ///
58    /// # Returns
59    ///
60    /// Returns an optimizer that can be used to optimize a module.
61    pub fn init(&self) -> ModuleOptimizer {
62        let mut optim = ModuleOptimizer::from(self.build());
63        if let Some(config) = &self.grad_clipping {
64            optim = optim.with_grad_clipping(config.init());
65        }
66
67        optim
68    }
69}
70
71/// Optimizer that implements stochastic gradient descent with momentum.
72/// The optimizer can be configured with [RmsPropConfig](RmsPropConfig).
73#[derive(Clone)]
74pub struct RmsProp {
75    alpha: f32,
76    // epsilon: f32,
77    centered: bool,
78    // momentum: Option<Momentum>,
79    momentum: RmsPropMomentum,
80    weight_decay: Option<WeightDecay>,
81}
82
83impl Optimizer for RmsProp {
84    type State<const D: usize> = RmsPropState<D>;
85
86    fn step<const D: usize>(
87        &self,
88        lr: LearningRate,
89        tensor: Tensor<D>,
90        mut grad: Tensor<D>,
91        state: Option<Self::State<D>>,
92    ) -> (Tensor<D>, Option<Self::State<D>>) {
93        // fetch state for params
94        let mut state_square_avg = None;
95        let mut state_centered = None;
96        let mut state_momentum = None;
97        if let Some(state) = state {
98            state_square_avg = Some(state.square_avg);
99            state_centered = Some(state.centered);
100            state_momentum = state.momentum;
101        }
102
103        // weight_decay transform
104        if let Some(weight_decay) = &self.weight_decay {
105            grad = weight_decay.transform(grad, tensor.clone());
106        }
107
108        // square_avg transform
109        let (grad, state_square_avg) =
110            SquareAvgState::transform(self.alpha, grad, state_square_avg);
111
112        // centered transform
113        let (grad, state_square_avg, state_centered) = CenteredState::transform(
114            self.alpha,
115            self.centered,
116            grad,
117            state_square_avg,
118            state_centered,
119        );
120
121        // momentum transform
122        let (grad, state_centered, state_momentum) =
123            self.momentum
124                .transform(grad, state_centered, state_momentum);
125
126        // transition state
127        let state = RmsPropState::new(state_square_avg, state_centered, state_momentum);
128
129        // tensor param transform
130        let delta = grad.mul_scalar(lr);
131        (tensor - delta, Some(state))
132    }
133
134    fn to_device<const D: usize>(mut state: Self::State<D>, device: &Device) -> Self::State<D> {
135        state.square_avg = state.square_avg.to_device(device);
136        state.centered = state.centered.to_device(device);
137        state.momentum = state.momentum.map(|momentum| momentum.to_device(device));
138        state
139    }
140}
141
142/// State of [RmsProp](RmsProp)
143#[derive(RecordState, Clone, new)]
144pub struct RmsPropState<const D: usize> {
145    /// Current squared average state.
146    pub square_avg: SquareAvgState<D>,
147    /// Current centered state
148    pub centered: CenteredState<D>,
149    /// Current gradient momentum, if any.
150    pub momentum: Option<RmsPropMomentumState<D>>,
151}
152
153/// [SquareAvgState](SquareAvgState) is to store and pass optimizer step params.
154#[derive(RecordState, Clone, new)]
155pub struct SquareAvgState<const D: usize> {
156    /// Current squared average.
157    pub square_avg: Tensor<D>,
158}
159
160impl<const D: usize> SquareAvgState<D> {
161    /// transform [SquareAvgState] to the next step
162    fn transform(alpha: f32, grad: Tensor<D>, state: Option<Self>) -> (Tensor<D>, Self) {
163        match state {
164            Some(state) => {
165                let square_avg = state
166                    .square_avg
167                    .mul_scalar(alpha)
168                    .add(grad.clone().square().mul_scalar(1. - alpha));
169                (grad, Self { square_avg })
170            }
171            _ => {
172                let square_avg = grad.clone().square().mul_scalar(1. - alpha);
173                (grad, Self { square_avg })
174            }
175        }
176    }
177
178    /// Moves the state to a device.
179    ///
180    /// # Arguments
181    ///
182    /// * `device` - Device to move the state to.
183    ///
184    /// # Returns
185    ///
186    /// * `self` - Moved state.
187    pub fn to_device(mut self, device: &Device) -> Self {
188        self.square_avg = self.square_avg.to_device(device);
189        self
190    }
191}
192
193/// [CenteredState](CenteredState) is to store and pass optimizer step params.
194#[derive(RecordState, Clone, new)]
195pub struct CenteredState<const D: usize> {
196    /// The averaged gradient to calculate the centered gradient, if available.
197    pub grad_avg: Option<Tensor<D>>,
198    /// The current average value.
199    pub avg: Tensor<D>,
200}
201
202impl<const D: usize> CenteredState<D> {
203    /// transform [CenteredState] to the next step
204    fn transform(
205        alpha: f32,
206        centered: bool,
207        grad: Tensor<D>,
208        square_avg_state: SquareAvgState<D>,
209        centered_state: Option<Self>,
210    ) -> (Tensor<D>, SquareAvgState<D>, Self) {
211        if centered {
212            let grad_avg_constant = grad.clone().mul_scalar(1. - alpha);
213            let grad_avg = match centered_state {
214                Some(state) => state
215                    .grad_avg
216                    .map_or(grad_avg_constant.clone(), move |grad_avg| {
217                        grad_avg.mul_scalar(alpha).add(grad_avg_constant)
218                    }),
219                _ => grad_avg_constant,
220            };
221            let avg = square_avg_state
222                .square_avg
223                .clone()
224                .sub(grad_avg.clone().square());
225
226            (
227                grad,
228                square_avg_state,
229                Self {
230                    grad_avg: Some(grad_avg),
231                    avg,
232                },
233            )
234        } else {
235            (
236                grad,
237                square_avg_state.clone(),
238                Self {
239                    grad_avg: None,
240                    avg: square_avg_state.square_avg,
241                },
242            )
243        }
244    }
245
246    /// Moves the state to a device.
247    ///
248    /// # Arguments
249    ///
250    /// * `device` - Device to move the state to.
251    ///
252    /// # Returns
253    ///
254    /// * `self` - Moved state.
255    pub fn to_device(mut self, device: &Device) -> Self {
256        self.grad_avg = self.grad_avg.map(|grad_avg| grad_avg.to_device(device));
257        self.avg = self.avg.to_device(device);
258        self
259    }
260}
261
262/// [RmsPropMomentum](RmsPropMomentum) is to store config status for optimizer.
263/// (, which is stored in [optimizer](RmsProp) itself and not passed in during `step()` calculation)
264#[derive(Clone)]
265pub struct RmsPropMomentum {
266    momentum: f32,
267    epsilon: f32,
268}
269
270impl RmsPropMomentum {
271    /// transform [grad](Tensor) and [RmsPropMomentumState] to the next step
272    fn transform<const D: usize>(
273        &self,
274        grad: Tensor<D>,
275        centered_state: CenteredState<D>,
276        momentum_state: Option<RmsPropMomentumState<D>>,
277    ) -> (Tensor<D>, CenteredState<D>, Option<RmsPropMomentumState<D>>) {
278        let grad = grad.div(centered_state.avg.clone().sqrt().add_scalar(self.epsilon));
279
280        if self.momentum > 0. {
281            let buf = match momentum_state {
282                Some(state) => state.buf.mul_scalar(self.momentum).add(grad),
283                _ => grad,
284            };
285            (
286                buf.clone(),
287                centered_state,
288                Some(RmsPropMomentumState { buf }),
289            )
290        } else {
291            (grad, centered_state, None)
292        }
293    }
294}
295
296/// [RmsPropMomentumState](RmsPropMomentumState) is to store and pass optimizer step params.
297#[derive(RecordState, Clone, new)]
298pub struct RmsPropMomentumState<const D: usize> {
299    buf: Tensor<D>,
300}
301
302impl<const D: usize> RmsPropMomentumState<D> {
303    /// Moves the state to a device.
304    ///
305    /// # Arguments
306    ///
307    /// * `device` - Device to move the state to.
308    ///
309    /// # Returns
310    ///
311    /// * `self` - Moved state.
312    pub fn to_device(mut self, device: &Device) -> Self {
313        self.buf = self.buf.to_device(device);
314        self
315    }
316}
317
318#[cfg(test)]
319mod tests {
320    use burn::tensor::Tolerance;
321
322    use super::*;
323    use crate::optim::GradientsParams;
324    use burn::module::Param;
325    use burn::tensor::{Distribution, Tensor, TensorData};
326    use burn_nn::{Linear, LinearConfig};
327
328    type FT = f32;
329
330    const LEARNING_RATE: LearningRate = 0.01;
331
332    #[test]
333    fn test_rmsprop_optimizer_save_load_state() {
334        let device = Device::default().autodiff();
335        let linear = LinearConfig::new(6, 6).init(&device);
336        let x = Tensor::<2>::random([2, 6], Distribution::Default, &device);
337        let mut optimizer = create_rmsprop();
338        let grads = linear.forward(x).backward();
339        let grads = GradientsParams::from_grads(grads, &linear);
340        let _linear = optimizer.step(LEARNING_RATE, linear, grads);
341
342        let bytes = optimizer.into_bytes().unwrap();
343        assert!(!bytes.is_empty());
344
345        #[cfg(feature = "std")]
346        optimizer
347            .save(std::env::temp_dir().as_path().join("test_optim_rmsprop"))
348            .unwrap();
349
350        let state_optim_before = optimizer.to_record();
351        let optimizer = create_rmsprop().from_bytes(bytes).unwrap();
352        let state_optim_after = optimizer.to_record();
353
354        assert_eq!(state_optim_before.len(), state_optim_after.len());
355    }
356
357    /// used for test differences and debug
358    #[test]
359    fn test_rmsprop_optimizer_with_numbers_basic() {
360        let device = Device::default().autodiff();
361        let linear = given_linear_layer(
362            TensorData::from([
363                [1., 1., 1., 1., 1., 1.],
364                [1., 1., 1., 1., 1., 1.],
365                [1., 1., 1., 1., 1., 1.],
366                [1., 1., 1., 1., 1., 1.],
367                [1., 1., 1., 1., 1., 1.],
368                [1., 1., 1., 1., 1., 1.],
369            ]),
370            TensorData::from([0.5, 0.5, 0.5, 0.5, 0.5, 0.5]),
371            &device,
372        );
373        let x_1 = Tensor::<2>::from_floats(
374            [
375                [0.6294, 0.0940, 0.8176, 0.8824, 0.5228, 0.4310],
376                [0.7152, 0.9559, 0.7893, 0.5684, 0.5939, 0.8883],
377            ],
378            &device,
379        )
380        .require_grad();
381        let x_2 = Tensor::<2>::from_floats(
382            [
383                [0.8491, 0.2108, 0.8939, 0.4433, 0.5527, 0.2528],
384                [0.3270, 0.0412, 0.5538, 0.9605, 0.3195, 0.9085],
385            ],
386            &device,
387        )
388        .require_grad();
389
390        let mut optimizer = RmsPropConfig::new()
391            .with_alpha(0.99)
392            .with_epsilon(1e-8)
393            .with_weight_decay(WeightDecayConfig::new(0.05).into())
394            .with_momentum(0.9)
395            .with_centered(false)
396            .init();
397
398        // println!("linear is {:?}", linear);
399        let grads = linear.forward(x_1).backward();
400        let grads = GradientsParams::from_grads(grads, &linear);
401        let linear = optimizer.step(LEARNING_RATE, linear, grads);
402
403        // println!("linear is {:?}", linear);
404        let grads = linear.forward(x_2).backward();
405        let grads = GradientsParams::from_grads(grads, &linear);
406        let linear = optimizer.step(LEARNING_RATE, linear, grads);
407
408        // println!("linear is {:?}", linear);
409        let state_updated = linear;
410
411        let (weight_updated, bias_updated) = (
412            state_updated.weight.to_data(),
413            state_updated.bias.unwrap().to_data(),
414        );
415
416        // println!("\nweight_updated\n{:?}", weight_updated);
417        // println!("\nbias_updated\n{:?}", bias_updated);
418
419        let weights_expected = TensorData::from([
420            [0.743937, 0.743937, 0.743937, 0.743937, 0.743937, 0.743937],
421            [0.783809, 0.783809, 0.783809, 0.783809, 0.783809, 0.783809],
422            [0.742881, 0.742881, 0.742881, 0.742881, 0.742881, 0.742881],
423            [0.740366, 0.740366, 0.740366, 0.740366, 0.740366, 0.740366],
424            [0.748005, 0.748005, 0.748005, 0.748005, 0.748005, 0.748005],
425            [0.743710, 0.743710, 0.743710, 0.743710, 0.743710, 0.743710],
426        ]);
427        let bias_expected =
428            TensorData::from([0.239199, 0.239199, 0.239199, 0.239199, 0.239199, 0.239199]);
429
430        let tolerance = Tolerance::absolute(1e-6);
431        bias_updated.assert_approx_eq::<FT>(&bias_expected, tolerance);
432        weight_updated.assert_approx_eq::<FT>(&weights_expected, tolerance);
433    }
434
435    #[test]
436    fn test_rmsprop_optimizer_with_numbers() {
437        let device = Device::default().autodiff();
438        let linear = given_linear_layer(
439            TensorData::from([
440                [-0.3206, 0.1374, 0.4043, 0.3200, 0.0859, 0.0671],
441                [0.0777, -0.0185, -0.3667, 0.2550, 0.1955, -0.2922],
442                [-0.0190, 0.0346, -0.2962, 0.2484, -0.2780, 0.3130],
443                [-0.2980, -0.2214, -0.3715, -0.2981, -0.0761, 0.1626],
444                [0.3300, -0.2182, 0.3717, -0.1729, 0.3796, -0.0304],
445                [-0.0159, -0.0120, 0.1258, 0.1921, 0.0293, 0.3833],
446            ]),
447            TensorData::from([-0.3905, 0.0884, -0.0970, 0.1176, 0.1366, 0.0130]),
448            &device,
449        );
450        let x_1 = Tensor::<2>::from_floats(
451            [
452                [0.6294, 0.0940, 0.8176, 0.8824, 0.5228, 0.4310],
453                [0.7152, 0.9559, 0.7893, 0.5684, 0.5939, 0.8883],
454            ],
455            &device,
456        )
457        .require_grad();
458        let x_2 = Tensor::<2>::from_floats(
459            [
460                [0.8491, 0.2108, 0.8939, 0.4433, 0.5527, 0.2528],
461                [0.3270, 0.0412, 0.5538, 0.9605, 0.3195, 0.9085],
462            ],
463            &device,
464        )
465        .require_grad();
466
467        let mut optimizer = RmsPropConfig::new()
468            .with_alpha(0.99)
469            .with_epsilon(1e-8)
470            .with_weight_decay(WeightDecayConfig::new(0.05).into())
471            .with_momentum(0.9)
472            .with_centered(false)
473            .init();
474
475        let grads = linear.forward(x_1).backward();
476        let grads = GradientsParams::from_grads(grads, &linear);
477        let linear = optimizer.step(LEARNING_RATE, linear, grads);
478
479        let grads = linear.forward(x_2).backward();
480        let grads = GradientsParams::from_grads(grads, &linear);
481        let linear = optimizer.step(LEARNING_RATE, linear, grads);
482
483        let state_updated = linear;
484        let weights_expected = TensorData::from([
485            [
486                -0.576399, -0.118494, 0.148353, 0.064070, -0.169983, -0.188779,
487            ],
488            [
489                -0.135571, -0.231448, -0.578445, 0.041143, -0.018162, -0.504207,
490            ],
491            [
492                -0.275990, -0.222397, -0.553153, -0.008625, -0.534956, 0.055967,
493            ],
494            [
495                -0.557575, -0.480979, -0.631072, -0.557675, -0.335686, -0.096997,
496            ],
497            [
498                0.078313, -0.469618, 0.119993, -0.424341, 0.127890, -0.281912,
499            ],
500            [
501                -0.271996, -0.268097, -0.130324, -0.064037, -0.226805, 0.127126,
502            ],
503        ]);
504        let bias_expected = TensorData::from([
505            -0.651299, -0.172400, -0.357800, -0.143200, -0.124200, -0.247800,
506        ]);
507
508        let (weight_updated, bias_updated) = (
509            state_updated.weight.to_data(),
510            state_updated.bias.unwrap().to_data(),
511        );
512
513        // println!("\nweight_updated\n{:?}", weight_updated);
514        // println!("\nbias_updated\n{:?}", bias_updated);
515
516        let tolerance = Tolerance::absolute(1e-6);
517        bias_updated.assert_approx_eq::<FT>(&bias_expected, tolerance);
518        weight_updated.assert_approx_eq::<FT>(&weights_expected, tolerance);
519    }
520
521    fn given_linear_layer(weight: TensorData, bias: TensorData, device: &Device) -> Linear {
522        Linear {
523            weight: Param::from_data(weight, device),
524            bias: Some(Param::from_data(bias, device)),
525        }
526    }
527
528    fn create_rmsprop() -> ModuleOptimizer {
529        RmsPropConfig {
530            alpha: 0.99,
531            epsilon: 1e-9,
532            centered: false,
533            weight_decay: Some(WeightDecayConfig { penalty: 0.05 }),
534            momentum: 0.9,
535            grad_clipping: None,
536        }
537        .init()
538    }
539}
burn_optim/optim/rmsprop.rs

burn_optim/optim/
rmsprop.rs
