burn 0.3.0

BURN: Burn Unstoppable Rusty Neurons
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

use burn::module::{Module, Param};
use burn::tensor::backend::Backend;
use burn::tensor::{Distribution, Shape, Tensor};

pub type TestBackend = burn_ndarray::NdArrayBackend<f32>;

#[derive(Module, Debug)]
struct ModuleBasic<B>
where
    B: Backend,
{
    weight_basic: Param<Tensor<B, 2>>,
}

impl<B: Backend> ModuleBasic<B> {
    fn new() -> Self {
        let weight_basic = Tensor::random(Shape::new([20, 20]), Distribution::Standard);
        Self {
            weight_basic: Param::new(weight_basic),
        }
    }
}

#[derive(Module, Debug)]
struct ModuleComposed<B>
where
    B: Backend,
{
    weight: Param<Tensor<B, 2>>,
    basic: Param<ModuleBasic<B>>,
}

impl<B: Backend> ModuleComposed<B> {
    fn new() -> Self {
        let weight = Tensor::random(Shape::new([20, 20]), Distribution::Standard);
        Self {
            weight: Param::new(weight),
            basic: Param::new(ModuleBasic::new()),
        }
    }
}

mod state {
    use super::*;

    #[test]
    fn should_load_from_state_basic() {
        let module_1 = ModuleBasic::<TestBackend>::new();
        let mut module_2 = ModuleBasic::<TestBackend>::new();
        let state_1 = module_1.state();
        assert_ne!(
            module_1.weight_basic.to_data(),
            module_2.weight_basic.to_data()
        );

        module_2.load(&state_1).unwrap();
        assert_eq!(
            module_1.weight_basic.to_data(),
            module_2.weight_basic.to_data()
        );
    }

    #[test]
    fn should_load_from_state_compose() {
        let module_1 = ModuleComposed::<TestBackend>::new();
        let mut module_2 = ModuleComposed::<TestBackend>::new();
        assert_ne!(module_1.weight.to_data(), module_2.weight.to_data());
        assert_ne!(
            module_1.basic.weight_basic.to_data(),
            module_2.basic.weight_basic.to_data()
        );

        let state_1 = module_1.state();
        module_2.load(&state_1).unwrap();

        assert_eq!(module_1.weight.to_data(), module_2.weight.to_data());
        assert_eq!(
            module_1.basic.weight_basic.to_data(),
            module_2.basic.weight_basic.to_data()
        );
    }
}

mod num_params {
    use super::*;

    #[test]
    fn should_calculate_num_params_basic() {
        let module = ModuleBasic::<TestBackend>::new();
        assert_eq!(20 * 20, module.num_params());
    }

    #[test]
    fn should_output_state_composed() {
        let module = ModuleComposed::<TestBackend>::new();
        assert_eq!(2 * 20 * 20, module.num_params());
    }
}