manopt_rs/optimizers/
many_steps.rs

1//! A optimizer that allows for many steps with a given learning schedule
2//! and a way of evaluating the gradient function on arbitrary
3//! points. This way we can step using `SimpleOptimizer::step` with that gradient
4//! several times.
5use crate::prelude::*;
6use burn::optim::{LearningRate, SimpleOptimizer};
7
8/// A optimizer that allows for many steps with a given learning schedule
9/// and a way of evaluating the gradient function on arbitrary
10/// points. This way we can step using `SimpleOptimizer::step` with that gradient
11/// several times.
12pub trait LessSimpleOptimizer<B: Backend>: SimpleOptimizer<B> {
13    fn many_steps<const D: usize>(
14        &self,
15        lr_function: impl FnMut(usize) -> LearningRate,
16        num_steps: usize,
17        grad_function: impl FnMut(Tensor<B, D>) -> Tensor<B, D>,
18        tensor: Tensor<B, D>,
19        state: Option<Self::State<D>>,
20    ) -> (Tensor<B, D>, Option<Self::State<D>>);
21}
22
23/// The implementation of `LessSimpleOptimizer` is completely determined
24/// by how `SimpleOptimizer` has been implemented because we are
25/// just taking gradients using the input `grad_function` and steping with
26/// `SimpleOptimizer::step`
27impl<B: Backend, T: SimpleOptimizer<B>> LessSimpleOptimizer<B> for T {
28    #[inline]
29    fn many_steps<const D: usize>(
30        &self,
31        mut lr_function: impl FnMut(usize) -> LearningRate,
32        num_steps: usize,
33        mut grad_function: impl FnMut(Tensor<B, D>) -> Tensor<B, D>,
34        mut tensor: Tensor<B, D>,
35        mut state: Option<Self::State<D>>,
36    ) -> (Tensor<B, D>, Option<Self::State<D>>) {
37        // Perform optimization steps
38        for step in 0..num_steps {
39            // Compute gradient at tensor
40            let cur_grad = grad_function(tensor.clone());
41            // The current learning rate for this step
42            let cur_lr = lr_function(step);
43            // Perform optimizer step
44            let (new_x, new_state) = self.step(cur_lr, tensor.clone(), cur_grad, state);
45            tensor = new_x.detach().require_grad();
46            state = new_state;
47        }
48        (tensor, state)
49    }
50}
manopt_rs/optimizers/many_steps.rs

manopt_rs/optimizers/
many_steps.rs
