mininn 0.1.4

A minimalist deep learnig crate for rust.
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

use mininn_derive::Layer;
use ndarray::{Array1, ArrayD, ArrayViewD};
use serde::{Deserialize, Serialize};

use crate::{
    core::{NNMode, NNResult},
    layers::{Layer, TrainLayer},
    utils::{MSGPackFormatting, Optimizer},
};

/// Flattens the input into a 1D array.
///
/// This layer is useful when the input is a 2D array, but you want to treat it as a 1D array.
///
/// ## Attributes
///
/// - `input`: The input array.
/// - `original_shape`: The original shape of the input array.
///
/// ## Example
///
/// ```
/// use ndarray::array;
/// use mininn::prelude::*;
///
/// let mut layer = Flatten::new();
/// let input = array![[1., 2., 3.], [4., 5., 6.]].into_dyn();
/// let output = layer.forward(input.view(), &NNMode::Train).unwrap();
/// assert_eq!(output.shape(), &[6]);
/// ```
#[derive(Layer, Debug, Clone, Serialize, Deserialize)]
pub struct Flatten {
    input: Array1<f32>,
    original_shape: Vec<usize>,
}

impl Flatten {
    /// Creates a new instance of the Flatten layer.
    pub fn new() -> Self {
        Self {
            input: Array1::zeros(0),
            original_shape: Vec::new(),
        }
    }
}

impl TrainLayer for Flatten {
    fn forward(&mut self, input: ArrayViewD<f32>, _mode: &NNMode) -> NNResult<ArrayD<f32>> {
        self.original_shape = input.shape().to_vec();
        self.input = input.flatten().to_owned();
        Ok(self.input.clone().into_dyn())
    }

    fn backward(
        &mut self,
        output_gradient: ArrayViewD<f32>,
        _learning_rate: f32,
        _optimizer: &Optimizer,
        _mode: &NNMode,
    ) -> NNResult<ArrayD<f32>> {
        let reshaped_gradient = output_gradient
            .to_shape(self.original_shape.clone())?
            .to_owned();
        Ok(reshaped_gradient.into_dyn())
    }
}

#[cfg(test)]
mod tests {
    use ndarray::array;

    use super::*;

    #[test]
    fn test_flatten_forward() {
        let mut layer = Flatten::new();
        let input = array![[1., 2., 3.], [4., 5., 6.]].into_dyn();
        let output = layer.forward(input.view(), &NNMode::Train).unwrap();
        assert_eq!(output.shape(), &[6]);
    }

    #[test]
    fn test_flatten_backward() {
        let mut layer = Flatten::new();
        let input = array![[1., 2., 3.], [4., 5., 6.]].into_dyn();
        let output = layer.forward(input.view(), &NNMode::Train).unwrap();

        assert_eq!(output.shape(), &[6]);

        let output_gradient = layer
            .backward(output.view(), 0.1, &Optimizer::GD, &NNMode::Train)
            .unwrap();

        assert_eq!(output_gradient.shape(), &[2, 3]);
    }
}