burn-core 0.22.0

Flexible and Comprehensive Deep Learning Framework in Rust
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
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134

//! Integration test for the non-generic record system (`burn_core::store`).
//!
//! Saves a (nested) module to a burnpack record, loads it back into a freshly initialized
//! module, and checks the parameters round-trip — exercising the public API exactly as a user
//! would (`into_record` / `save` / `load` / `load_record`).

use burn::module::{Module, Param};
use burn::store::ModuleRecord;
use burn::tensor::Tensor;
use burn_core as burn;
use burn_tensor::Device;

#[derive(Module, Debug)]
struct Layer {
    weight: Param<Tensor<2>>,
    bias: Param<Tensor<1>>,
}

impl Layer {
    fn from_values(weight: [[f32; 2]; 2], bias: [f32; 2], device: &Device) -> Self {
        Self {
            weight: Param::from_data(weight, device),
            bias: Param::from_data(bias, device),
        }
    }

    fn zeros(device: &Device) -> Self {
        Self::from_values([[0.0; 2]; 2], [0.0; 2], device)
    }

    fn values(&self) -> (Vec<f32>, Vec<f32>) {
        (
            self.weight.val().try_into_vec_as().unwrap(),
            self.bias.val().try_into_vec_as().unwrap(),
        )
    }
}

/// A module with only the `first` layer — its record holds `first.weight` / `first.bias`,
/// a strict subset of [`Mlp`]'s parameters.
#[derive(Module, Debug)]
struct FirstOnly {
    first: Layer,
}

/// Nested module so the record exercises path building (`first.weight`, `second.bias`, ...).
#[derive(Module, Debug)]
struct Mlp {
    first: Layer,
    second: Layer,
}

impl Mlp {
    fn sample(device: &Device) -> Self {
        Self {
            first: Layer::from_values([[1.0, 2.0], [3.0, 4.0]], [5.0, 6.0], device),
            second: Layer::from_values([[7.0, 8.0], [9.0, 10.0]], [11.0, 12.0], device),
        }
    }

    fn zeros(device: &Device) -> Self {
        Self {
            first: Layer::zeros(device),
            second: Layer::zeros(device),
        }
    }
}

fn assert_matches_sample(model: &Mlp) {
    let (w1, b1) = model.first.values();
    let (w2, b2) = model.second.values();
    assert_eq!(w1, vec![1.0, 2.0, 3.0, 4.0]);
    assert_eq!(b1, vec![5.0, 6.0]);
    assert_eq!(w2, vec![7.0, 8.0, 9.0, 10.0]);
    assert_eq!(b2, vec![11.0, 12.0]);
}

#[cfg(feature = "std")]
#[test]
fn save_and_load_module_via_file() {
    let device = Default::default();
    let dir = tempfile::tempdir().unwrap();
    let path = dir.path().join("mlp.bpk");

    // Save a trained module to a record on disk.
    let record = Mlp::sample(&device).into_record();
    assert_eq!(record.len(), 4); // first.weight, first.bias, second.weight, second.bias
    record.save(&path).unwrap();

    // Load the record and apply it onto a freshly initialized (zeroed) module.
    let record = ModuleRecord::load(&path).unwrap();
    let loaded = Mlp::zeros(&device).load_record(record);

    assert_matches_sample(&loaded);
}

#[test]
fn save_and_load_module_via_bytes() {
    let device = Default::default();

    let bytes = Mlp::sample(&device).into_record().into_bytes().unwrap();

    let record = ModuleRecord::from_bytes(bytes).unwrap();
    let loaded = Mlp::zeros(&device).load_record(record);

    assert_matches_sample(&loaded);
}

#[test]
fn missing_parameters_require_allow_partial() {
    let device = Default::default();

    // A record holding only `first.weight` / `first.bias` (a subset of the full Mlp).
    let partial = FirstOnly {
        first: Layer::from_values([[1.0, 2.0], [3.0, 4.0]], [5.0, 6.0], &device),
    }
    .into_record();
    let bytes = partial.into_bytes().unwrap();

    // Strict load into the full Mlp fails: `second.*` is missing from the record.
    let strict = ModuleRecord::from_bytes(bytes.clone()).unwrap();
    assert!(Mlp::zeros(&device).try_load_record(strict).is_err());

    // With `allow_partial`, the present params load and the rest keep their init.
    let lenient = ModuleRecord::from_bytes(bytes).unwrap().allow_partial(true);
    let loaded = Mlp::zeros(&device).load_record(lenient);

    let (w1, b1) = loaded.first.values();
    assert_eq!(w1, vec![1.0, 2.0, 3.0, 4.0]);
    assert_eq!(b1, vec![5.0, 6.0]);
    let (w2, b2) = loaded.second.values();
    assert_eq!(w2, vec![0.0, 0.0, 0.0, 0.0]);
    assert_eq!(b2, vec![0.0, 0.0]);
}