koho 0.1.0

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

use std::collections::HashMap;

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

use crate::nn::optim::Optimizer;

/// Stochastic Gradient Descent optimizer
pub struct SGD {
    lr: f64,
    momentum: Option<f64>,
    weight_decay: Option<f64>,
    momentum_buffers: HashMap<*const Var, Tensor>,
}

impl SGD {
    /// Create a new SGD optimizer
    pub fn new(lr: f64) -> CandleResult<Self> {
        Ok(Self {
            lr,
            momentum: None,
            weight_decay: None,
            momentum_buffers: HashMap::new(),
        })
    }

    /// Enable momentum with the given coefficient
    pub fn momentum(mut self, momentum: f64) -> Self {
        self.momentum = Some(momentum);
        self
    }

    /// Enable weight decay (L2 regularization)
    pub fn weight_decay(mut self, weight_decay: f64) -> Self {
        self.weight_decay = Some(weight_decay);
        self
    }
}

impl Optimizer for SGD {
    fn step(&mut self, grads: &GradStore, vars: Vec<&mut Var>) -> CandleResult<()> {
        for var in vars {
            let ptr = var as *const Var;

            if let Some(grad) = grads.get(var) {
                let mut effective_grad = grad.clone();

                if let Some(wd) = self.weight_decay {
                    effective_grad = (effective_grad + (var.as_tensor() * wd)?)?;
                }

                if let Some(momentum) = self.momentum {
                    let buf = if let Some(buf) = self.momentum_buffers.get(&ptr) {
                        (buf * momentum) + effective_grad.clone()
                    } else {
                        Ok(effective_grad.clone())
                    }?;

                    self.momentum_buffers.insert(ptr, buf.clone());
                    effective_grad = buf;
                }

                var.set(&(var.as_tensor() - (effective_grad * self.lr)?)?)?;
            }
        }
        Ok(())
    }

    fn learning_rate(&self) -> f64 {
        self.lr
    }

    fn set_learning_rate(&mut self, lr: f64) {
        self.lr = lr;
    }
}