1 files changed, 59 insertions, 0 deletions
diff --git a/src/model.rs b/src/model.rs
new file mode 100644
index 0000000..04eb6c8
--- /dev/null
+++ b/src/model.rs
@@ -0,0 +1,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)
+}