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
|
use burn::{
nn::conv::{Conv2d, Conv2dConfig},
nn::{Linear, Relu},
prelude::*,
};
use nn::{LeakyReluConfig, LinearConfig};
use burn::record::{FullPrecisionSettings, NamedMpkFileRecorder, Recorder};
use burn_import::pytorch::PyTorchFileRecorder;
type ArrayBackend = burn_ndarray::NdArray<f32>;
#[derive(Module, Debug)]
pub struct Net<B: Backend> {
input_lin: Linear<B>,
relu1: Relu,
lin2: Linear<B>,
relu2: Relu,
lin3: Linear<B>,
}
impl<B: Backend> Net<B> {
/// Create a new model.
pub fn init(device: &B::Device) -> Self {
let input_size = 8;
let input_lin = LinearConfig::new(input_size, 16).init(device);
let relu1 = Relu::new();
let lin2 = LinearConfig::new(16, 4).init(device);
let relu2 = Relu::new();
let lin3 = LinearConfig::new(4, 1).init(device);
Self {
input_lin,
relu1,
lin2,
relu2,
lin3,
}
}
/// Forward pass of the model.
pub fn forward(&self, x: Tensor<B, 1>) -> Tensor<B, 1> {
let x = self.input_lin.forward(x);
let x = self.relu1.forward(x);
let x = self.lin2.forward(x);
let x = self.relu2.forward(x);
let x = self.lin3.forward(x);
x
}
}
/// Load the model from the file in your source code (not in build.rs or script).
pub fn load_model() -> Net<ArrayBackend> {
let device = Default::default();
let record: NetRecord<ArrayBackend> = PyTorchFileRecorder::<FullPrecisionSettings>::default()
.load("./perf.pt".into(), &device)
.expect("Failed to decode state");
Net::<ArrayBackend>::init(&device).load_record(record)
}
|
