mushin 0.5.0

Computational graphs with reverse automatic differentation in the GPU
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

use crate::graph::node::Node;
use std::rc::Rc;

/// Stochastic Gradient Descent
pub struct SGD {
    lr: f32,
    params: Vec<Rc<Node>>,
}

impl SGD {
    #[inline]
    pub fn new<'n, P>(params: &'n P, lr: f32) -> Self
    where
        &'n P: IntoIterator<Item = &'n Rc<Node>>,
    {
        Self {
            lr,
            params: params
                .into_iter()
                .filter_map(|n| {
                    if n.is_declaration() {
                        Some(n.clone())
                    } else {
                        None
                    }
                })
                .collect(),
        }
    }

    #[inline]
    pub fn step(&self) {
        for node in &self.params {
            *node.data_mut() -= self.lr * &node.grad().clone();
        }
    }
}

#[cfg(test)]
mod tests {
    use super::SGD;
    use crate as mu;
    use crate::tests::equal_arrays;
    use crate::Tensor;

    #[test]
    fn sgd_step() {
        let x = mu::fill::<1, 1, 1, 1>(1.0);
        let optim = SGD::new(&[(&x).into()], 0.1);

        x.backward();
        optim.step();
        assert!(equal_arrays(x.data(), arrayfire::constant!(0.9; 1,1,1,1)));
    }
}