meuron 0.4.0

Meuron is a modular neural network library written in rust for training simple neural networks.
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

use crate::activation::Activation;
use crate::backend::{Backend, DefaultBackend};
use crate::initializer::Initializer;
use crate::layer::Layer;
use crate::optimizer::Optimizer;
use ndarray::{Ix1, Ix2};
use serde::{Deserialize, Serialize};
use std::marker::PhantomData;

#[derive(Serialize, Deserialize)]
#[serde(bound(
    serialize = "A: Serialize, B::Tensor<Ix2>: Serialize, B::Tensor<Ix1>: Serialize",
    deserialize = "A: Deserialize<'de>, B::Tensor<Ix2>: Deserialize<'de>, B::Tensor<Ix1>: Deserialize<'de>"
))]
pub struct DenseLayer<A, B: Backend = DefaultBackend>
where
    A: Activation<B>,
{
    pub weights: B::Tensor<Ix2>,
    pub biases: B::Tensor<Ix1>,
    pub activation: A,
    #[serde(skip)]
    last_input: Option<B::Tensor<Ix2>>,
    #[serde(skip)]
    last_z: Option<B::Tensor<Ix2>>,
    #[serde(skip)]
    grad_weights: Option<B::Tensor<Ix2>>,
    #[serde(skip)]
    grad_biases: Option<B::Tensor<Ix1>>,
    #[serde(skip)]
    _backend: PhantomData<B>,
}

impl<A: Activation<B>, B: Backend> DenseLayer<A, B> {
    pub fn new<IW, IB>(
        input_size: usize,
        output_size: usize,
        activation: A,
        weight_init: IW,
        bias_init: IB,
    ) -> Self
    where
        IW: Initializer<B>,
        IB: Initializer<B>,
    {
        DenseLayer {
            weights: weight_init.init(Ix2(input_size, output_size)),
            biases: bias_init.init(Ix1(output_size)),
            activation,
            last_input: None,
            last_z: None,
            grad_weights: None,
            grad_biases: None,
            _backend: PhantomData,
        }
    }
}

impl<A: Activation<B>, B: Backend> Layer<B> for DenseLayer<A, B> {
    type Input = Ix2;
    type Output = Ix2;

    fn forward(&mut self, input: &B::Tensor<Ix2>) -> B::Tensor<Ix2> {
        self.last_input = Some(input.clone());
        let z = B::broadcast_add(&B::matmul(input, &self.weights), &self.biases);
        self.last_z = Some(z.clone());
        self.activation.activate(&z)
    }

    fn backward(&mut self, grad_output: &B::Tensor<Ix2>) -> B::Tensor<Ix2> {
        let last_z = self.last_z.as_ref().expect("forward before backward");
        let last_input = self.last_input.as_ref().expect("forward before backward");

        let grad_z = self.activation.vjp(last_z, grad_output);

        let batch_size = B::len_of(last_input, 0) as f32;
        let inv_batch = 1.0 / batch_size.max(1.0);

        self.grad_weights = Some(B::scale(
            &B::matmul(&B::transpose(last_input, 0, 1), &grad_z),
            inv_batch,
        ));
        self.grad_biases = Some(B::scale(&B::sum_axis(&grad_z, 0), inv_batch));

        B::matmul(&grad_z, &B::transpose(&self.weights, 0, 1))
    }

    fn update<O: Optimizer<B>>(&mut self, optimizer: &mut O) {
        if let (Some(gw), Some(gb)) = (self.grad_weights.take(), self.grad_biases.take()) {
            optimizer.update_param(&mut self.weights, &gw);
            optimizer.update_param(&mut self.biases, &gb);
        }
    }
}