yarnn 0.1.0

Yet Another rust Neural Network framework
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

use crate::tensor::{Tensor, TensorShape};
use crate::layer::Layer;
use crate::params::Params;
use crate::backend::{Backend, BackendGemm, BackendBias, BackendScale};
use crate::optimizer::{Optimizable, Optimizer};

pub struct LinearConfig {
    pub units: u32,
    pub biases: bool,
}

impl Default for LinearConfig {
    fn default() -> Self {
        Self {
            units: 1,
            biases: false,
        }
    }
}

pub struct Linear<N, B, O> 
    where B: Backend<N>,
          O: Optimizer<N, B>
{
    inputs: u32,
    outputs: u32,
    use_biases: bool,
    weights: Params<N, B, O>,
    biases: Params<N, B, O>,
}

impl <N, B, O> Layer<N, B> for Linear<N, B, O> 
    where B: Backend<N> + BackendGemm<N> + BackendBias<N>,
          O: Optimizer<N, B>
{
    type Config = LinearConfig;

    fn name(&self) -> &str {
        "Linear"
    }
    
    fn create(input_shape: TensorShape, cfg: Self::Config) -> Self {
        assert!(input_shape.dims == 1);

        let inputs = input_shape.get(0);

        Linear {
            inputs,
            outputs: cfg.units,
            use_biases: cfg.biases,
            weights: Params::new((inputs, cfg.units)),
            biases: Params::new((cfg.units, )),
        }
    }

    fn init(&mut self, backend: &B) {
        self.weights.init_random(backend, self.inputs + self.outputs);
        if self.use_biases {
            self.biases.init_zero(backend);
        }
    }

    #[inline]
    fn input_shape(&self) -> TensorShape {
        TensorShape::new1d(self.inputs)
    }

    #[inline]
    fn output_shape(&self) -> TensorShape {
        TensorShape::new1d(self.outputs)
    }
    
    #[inline]
    fn forward(&self, backend: &B, y: &mut B::Tensor, x: &B::Tensor) {
        backend.matmul(y, x, &self.weights.params);

        if self.use_biases {
            backend.bias_add(y, &self.biases.params);
        }
    }

    #[inline]
    fn backward(&self, backend: &B, dx: &mut B::Tensor, dy: &B::Tensor, _: &B::Tensor, _: &B::Tensor) {
        backend.matmul_nt(dx, dy, &self.weights.params);
    }
}

impl <N, B, O> Optimizable<N, B, O> for Linear<N, B, O>
    where B: Backend<N> + BackendGemm<N> + BackendBias<N> + BackendScale<N>,
          O: Optimizer<N, B>
{
    fn calc_gradients(&mut self, backend: &B, inputs: &B::Tensor, deltas: &B::Tensor) {
        let prescaler = 1.0 / inputs.shape().get(0) as f32;

        backend.matmul_tn(&mut self.weights.grads, inputs, deltas);
        backend.scale(&mut self.weights.grads, backend.scalar_f32(prescaler));

        if self.use_biases {
            backend.scale(&mut self.biases.grads, backend.scalar_f32(prescaler));
            backend.bias_grad(&mut self.biases.grads, deltas);
        }
    }

    #[inline]
    fn optimize(&mut self, backend: &B, optimizer: &O) {
        optimizer.update_params(backend, &mut self.weights.ctx, &mut self.weights.params, &mut self.weights.grads);

        if self.use_biases {
            optimizer.update_params(backend, &mut self.biases.ctx, &mut self.biases.params, &mut self.biases.grads);
        }
    }
}