manopt_rs/optimizers/
hessian_optimizer.rs

1use burn::{
2    optim::SimpleOptimizer, prelude::Backend, record::Record, tensor::Tensor, optim::LearningRate,
3};
4
5/// TODO document and construct an implementation
6pub trait SimpleHessianOptimizer<B>: Send + Sync + Clone + SimpleOptimizer<B>
7where
8    B: Backend,
9{
10    /// The state of the optimizer. It also implements [record](Record), so that it can be saved.
11    type StateWithHessian<const D: usize, const H: usize>: Record<B>
12        + Clone
13        + From<Self::State<D>>
14        + Into<Self::State<D>>
15        + 'static;
16
17    /// The optimizer step is performed for one tensor at a time with its gradient, hessian and state.
18    ///
19    /// Note that the state is passed as parameter, so implementations don't have to handle
20    /// the saving and loading of recorded states.
21    fn step_with_hessian<const D: usize, const H: usize>(
22        &self,
23        lr: LearningRate,
24        tensor: Tensor<B, D>,
25        grad: Tensor<B, D>,
26        hessian: Tensor<B, H>,
27        state: Option<Self::StateWithHessian<D, H>>,
28    ) -> (Tensor<B, D>, Option<Self::StateWithHessian<D, H>>);
29
30    /// The optimizer step is performed for one tensor at a time with its gradient, hessian and state.
31    ///
32    /// Note that the state is passed as parameter, so implementations don't have to handle
33    /// the saving and loading of recorded states.
34    fn step<const D: usize, const H: usize>(
35        &self,
36        lr: LearningRate,
37        tensor: Tensor<B, D>,
38        grad: Tensor<B, D>,
39        hessian: Option<Tensor<B, H>>,
40        state: Option<Self::StateWithHessian<D, H>>,
41    ) -> (Tensor<B, D>, Option<Self::StateWithHessian<D, H>>) {
42        if let Some(hessian) = hessian {
43            self.step_with_hessian(lr, tensor, grad, hessian, state)
44        } else {
45            let (new_pt, new_state) =
46                SimpleOptimizer::step(self, lr, tensor, grad, state.map(Into::into));
47            (new_pt, new_state.map(Into::into))
48        }
49    }
50
51    /// Change the device of the state.
52    ///
53    /// This function will be called accordindly to have the state on the same device as the
54    /// gradient and the tensor when the [step](SimpleOptimizer::step) function is called.
55    fn to_device<const D: usize, const H: usize>(
56        state: Self::StateWithHessian<D, H>,
57        device: &B::Device,
58    ) -> Self::StateWithHessian<D, H>;
59}
60
61/// A `SimpleOptimizer` also works as a `SimpleHessianOptimizer` by ignoring the Hessian information
62impl<B: Backend, T: SimpleOptimizer<B>> SimpleHessianOptimizer<B> for T {
63    type StateWithHessian<const D: usize, const H: usize> = <T as SimpleOptimizer<B>>::State<D>;
64
65    fn step_with_hessian<const D: usize, const H: usize>(
66        &self,
67        lr: LearningRate,
68        tensor: Tensor<B, D>,
69        grad: Tensor<B, D>,
70        _hessian: Tensor<B, H>,
71        state: Option<Self::StateWithHessian<D, H>>,
72    ) -> (Tensor<B, D>, Option<Self::StateWithHessian<D, H>>) {
73        self.step(lr, tensor, grad, state)
74    }
75
76    fn to_device<const D: usize, const H: usize>(
77        state: Self::StateWithHessian<D, H>,
78        device: &<B as Backend>::Device,
79    ) -> Self::StateWithHessian<D, H> {
80        <T as SimpleOptimizer<B>>::to_device(state, device)
81    }
82}
83
84/// A optimizer that allows for many steps with a given learning schedule
85/// and a way of evaluating the gradient and hessian functions on arbitrary
86/// points. This way we can step using `SimpleHessianOptimizer::step` with
87/// that gradient and hessian several times.
88pub trait LessSimpleHessianOptimizer<B: Backend>: SimpleHessianOptimizer<B> {
89    fn many_steps<const D: usize, const H: usize>(
90        &self,
91        lr_function: impl FnMut(usize) -> LearningRate,
92        num_steps: usize,
93        grad_function: impl FnMut(Tensor<B, D>) -> Tensor<B, D>,
94        hessian_function: impl FnMut(Tensor<B, D>) -> Option<Tensor<B, H>>,
95        tensor: Tensor<B, D>,
96        state: Option<Self::StateWithHessian<D, H>>,
97    ) -> (Tensor<B, D>, Option<Self::StateWithHessian<D, H>>);
98}
99
100/// The implementation of `LessSimpleHessianOptimizer` is completely determined
101/// by how `SimpleHessianOptimizer` has been implemented because we are
102/// just taking gradients using the input `grad_function` and hessians with `hessian_function`
103/// and steping with `SimpleHessianOptimizer::step`
104impl<B: Backend, T: SimpleHessianOptimizer<B>> LessSimpleHessianOptimizer<B> for T {
105    #[inline]
106    fn many_steps<const D: usize, const H: usize>(
107        &self,
108        mut lr_function: impl FnMut(usize) -> LearningRate,
109        num_steps: usize,
110        mut grad_function: impl FnMut(Tensor<B, D>) -> Tensor<B, D>,
111        mut hessian_function: impl FnMut(Tensor<B, D>) -> Option<Tensor<B, H>>,
112        mut tensor: Tensor<B, D>,
113        mut state: Option<Self::StateWithHessian<D, H>>,
114    ) -> (Tensor<B, D>, Option<Self::StateWithHessian<D, H>>) {
115        // Perform optimization steps
116        for step in 0..num_steps {
117            // Compute gradient at tensor
118            let cur_grad = grad_function(tensor.clone());
119            // Compute hessian at tensor
120            let cur_hessian = hessian_function(tensor.clone());
121            // The current learning rate for this step
122            let cur_lr = lr_function(step);
123            // Perform optimizer step
124            let (new_x, new_state) = SimpleHessianOptimizer::step(
125                self,
126                cur_lr,
127                tensor.clone(),
128                cur_grad,
129                cur_hessian,
130                state,
131            );
132            tensor = new_x.detach().require_grad();
133            state = new_state;
134        }
135        (tensor, state)
136    }
137}
manopt_rs/optimizers/hessian_optimizer.rs

manopt_rs/optimizers/
hessian_optimizer.rs
