zenu-layer 0.1.1

A simple neural network layer library.
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

use std::collections::HashMap;

use zenu_autograd::{
    creator::{ones::ones, zeros::zeros},
    nn::batch_norm::{batch_norm_2d, BatchNorm2dAutoGradConfig},
    Variable,
};
use zenu_matrix::{device::Device, dim::DimTrait, num::Num};

use crate::{Module, Parameters};

pub struct BatchNorm2d<T: Num, D: Device> {
    config: BatchNorm2dAutoGradConfig<T>,
    momentum: f64,
    pub scale: Variable<T, D>,
    pub bias: Variable<T, D>,
    pub mean: Variable<T, D>,
    pub variance: Variable<T, D>,
}

impl<T: Num, D: Device> Module<T, D> for BatchNorm2d<T, D> {
    type Input = Variable<T, D>;
    type Output = Variable<T, D>;
    fn call(&self, input: Variable<T, D>) -> Variable<T, D> {
        if input.get_shape() != self.config.get_shape() {
            self.config.update_shape(input.get_shape().slice());
        }
        batch_norm_2d(
            input,
            self.scale.clone(),
            self.bias.clone(),
            self.mean.clone(),
            self.variance.clone(),
            self.momentum,
            self.config.clone(),
        )
    }
}

impl<T: Num, D: Device> Parameters<T, D> for BatchNorm2d<T, D> {
    fn weights(&self) -> HashMap<String, Variable<T, D>> {
        let mut weights = HashMap::new();
        weights.insert("batch_norm_2d.scale".to_string(), self.scale.clone());
        weights
    }

    fn biases(&self) -> HashMap<String, Variable<T, D>> {
        let mut biases = HashMap::new();
        biases.insert("batch_norm_2d.bias".to_string(), self.bias.clone());
        biases
    }

    fn parameters(&self) -> HashMap<String, Variable<T, D>> {
        let mut parameters = HashMap::new();
        for (key, value) in &self.weights() {
            parameters.insert(key.clone(), value.clone());
        }
        for (key, value) in &self.biases() {
            parameters.insert(key.clone(), value.clone());
        }
        parameters.insert("batch_norm_2d.mean".to_string(), self.mean.clone());
        parameters.insert("batch_norm_2d.variance".to_string(), self.variance.clone());
        parameters
    }
}

impl<T: Num, D: Device> BatchNorm2d<T, D> {
    #[must_use]
    pub fn new(channels: usize, momentum: f64) -> Self {
        let scale = ones([channels]);
        let bias = zeros([channels]);
        let mean = zeros([channels]);
        let variance = ones([channels]);

        scale.set_is_train(true);
        bias.set_is_train(true);

        scale.set_name("batch_norm_2d.scale");
        bias.set_name("batch_norm_2d.bias");
        mean.set_name("batch_norm_2d.mean");
        variance.set_name("batch_norm_2d.variance");

        let config = BatchNorm2dAutoGradConfig::default();
        Self {
            config,
            momentum,
            scale,
            bias,
            mean,
            variance,
        }
    }

    pub fn to<Dout: Device>(self) -> BatchNorm2d<T, Dout> {
        BatchNorm2d {
            config: self.config,
            momentum: self.momentum,
            scale: self.scale.to(),
            bias: self.bias.to(),
            mean: self.mean.to(),
            variance: self.variance.to(),
        }
    }
}