luma-optim 0.1.0

Optimizer for update grads
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22

mod adam;
mod momentum;
mod rms_prop;
mod sgd;

use luma_io::lpk::LumaPack;
use luma_tensor::{Device, GradStore};

pub use adam::*;
pub use momentum::*;
pub use rms_prop::*;
pub use sgd::*;

pub trait Optimizer {
    type Device: Device;

    fn get_lr(&self) -> f64;
    fn set_lr(&mut self, lr: f64);
    fn step(&mut self, grads: &GradStore<Self::Device>) -> luma_tensor::Result<()>;
    fn state_dict(&self) -> luma_tensor::Result<LumaPack<Self::Device>>;
    fn load_state_dict(&mut self, pack: &LumaPack<Self::Device>) -> luma_tensor::Result<()>;
}