koho 0.1.1

A deep learning model for spectral diffusion over k-cells in arbitrary cell complexes, built on `candle`
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121

use candle_core::{backprop::GradStore, Result as CandleResult, Var};

pub use crate::nn::optim::adam::Adam;
pub use crate::nn::optim::rmsprop::RMSprop;
pub use crate::nn::optim::sgd::SGD;

mod adam;
mod rmsprop;
mod sgd;

#[derive(Clone, Copy, Debug, PartialEq, PartialOrd)]
pub enum Decay {
    /// Weight decay regularisation to penalise large weights
    WeightDecay(f64),
    /// Decoupled weight decay as described in [Decoupled Weight Decay Regularization](https://arxiv.org/abs/1711.05101)
    DecoupledWeightDecay(f64),
}

/// Type of momentum to use
#[derive(Copy, Clone, Debug, PartialEq, PartialOrd)]
pub enum Momentum {
    /// classical momentum
    Classical(f64),
    /// nesterov momentum
    Nesterov(f64),
}

/// Trait for optimizers that work with mutable references to variables
pub trait Optimizer {
    /// Perform one optimization step using the gradients
    fn step(&mut self, grads: &GradStore, vars: Vec<&mut Var>) -> CandleResult<()>;

    /// Get the current learning rate
    fn learning_rate(&self) -> f64;

    /// Set the learning rate (for scheduling)
    fn set_learning_rate(&mut self, lr: f64);
}

/// Extra parameters passed as input during to the `Optimize` creation for the network.
pub enum OptimizerParams {
    SGD {
        dampening: Option<f64>,
        weight_decay: Option<Decay>,
        momentum: Option<Momentum>,
    },
    AdamW {
        beta_1: f64,
        beta_2: f64,
        eps: f64,
        weight_decay: Option<Decay>,
        amsgrad: bool,
    },
    Else,
}

/// Enum to represent which optimizer type to use
#[derive(Debug, Clone)]
pub enum OptimKind {
    SGD,
    Adam,
    AdamW, // will use Adam with decoupled weight decay
    RMSprop,
}

// Helper function to create optimizers from your existing enum
pub fn create_optimizer<'a>(
    kind: OptimKind,
    vars: Vec<&mut Var>,
    lr: f64,
    params: OptimizerParams,
) -> CandleResult<Box<dyn Optimizer + 'a>> {
    match kind {
        OptimKind::SGD => {
            if let OptimizerParams::SGD {
                momentum,
                weight_decay,
                ..
            } = params
            {
                let mut opt = SGD::new(lr)?;
                if let Some(Momentum::Nesterov(val)) = momentum {
                    opt = opt.momentum(val);
                }
                if let Some(Decay::WeightDecay(val)) = weight_decay {
                    opt = opt.weight_decay(val);
                }
                Ok(Box::new(opt))
            } else {
                Ok(Box::new(SGD::new(lr)?))
            }
        }
        OptimKind::Adam | OptimKind::AdamW => {
            if let OptimizerParams::AdamW {
                beta_1,
                beta_2,
                eps,
                weight_decay,
                ..
            } = params
            {
                let mut opt = Adam::new(vars, lr)?.betas(beta_1, beta_2).eps(eps);

                if let Some(wd) = weight_decay {
                    match wd {
                        Decay::WeightDecay(val) => {
                            opt = opt.weight_decay(val);
                        }
                        Decay::DecoupledWeightDecay(val) => {
                            opt = opt.decoupled_weight_decay(val);
                        }
                    }
                }
                Ok(Box::new(opt))
            } else {
                Ok(Box::new(Adam::new(vars, lr)?))
            }
        }
        OptimKind::RMSprop => Ok(Box::new(RMSprop::new(vars, lr)?)),
    }
}