Implement MNIST model and inference

Signed-off-by: Aisuko <[email protected]>

.gitignore

@@ -43,3 +43,4 @@ prepare
 /ggml-metal.metal
 target/
 Cargo.lock
+model.bin

backend/rust/models/Cargo.toml

@@ -5,6 +5,12 @@ edition = "2021"
 
 # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
 
+[features]
+default = ["ndarray"]
+
+ndarray = ["burn/ndarray"]
+wgpu = ["burn/wgpu"]
+
 [dependencies]
-burn = { version="0.10.0", features=["ndarray"] } # https://github.com/mudler/LocalAI/discussions/1219
+burn = { version="0.10.0", features=["ndarray","wgpu"] } # https://github.com/mudler/LocalAI/discussions/1219
 serde = "1.0.190"

backend/rust/models/src/lib.rs

@@ -1 +1 @@
-pub(crate) mod onnx;
+pub(crate) mod mnist;

backend/rust/models/src/mnist/mod.rs

@@ -0,0 +1,185 @@
+//! Defination of a mninst model and config of it.
+//! The source code is from https://github.com/burn-rs/burn/blob/main/examples/mnist-inference-web/src/model.rs
+//! The license is Apache-2.0 and MIT.
+//! Adapter by Aisuko
+
+use burn::{
+    backend::wgpu::{compute::init_async, AutoGraphicsApi, WgpuDevice},
+    module::Module,
+    nn::{self, BatchNorm, PaddingConfig2d},
+    record::{BinBytesRecorder, FullPrecisionSettings, Recorder},
+    tensor::{backend::Backend, Tensor},
+};
+
+// https://github.com/burn-rs/burn/blob/main/examples/mnist-inference-web/model.bin
+static STATE_ENCODED: &[u8] = include_bytes!("model.bin");
+
+const NUM_CLASSES: usize = 10;
+
+#[derive(Module, Debug)]
+/// A struct representing an MNINST model.
+pub struct Model<B: Backend> {
+    /// The first convolutional block of the model.
+    conv1: ConvBlock<B>,
+    /// The second convolutional block of the model.
+    conv2: ConvBlock<B>,
+    /// The third convolutional block of the model.
+    conv3: ConvBlock<B>,
+    /// A dropout layer used in the model.
+    dropout: nn::Dropout,
+    /// The first fully connected layer of the model.
+    fc1: nn::Linear<B>,
+    /// The second fully connected layer of the model.
+    fc2: nn::Linear<B>,
+    /// The activation function used in the model.
+    activation: nn::GELU,
+}
+
+impl<B: Backend> Model<B> {
+    pub fn new() -> Self {
+        let conv1 = ConvBlock::new([1, 8], [3, 3]); // 1 input channel, 8 output channels, 3x3 kernel size
+        let conv2 = ConvBlock::new([8, 16], [3, 3]); // 8 input channels, 16 output channels, 3x3 kernel size
+        let conv3 = ConvBlock::new([16, 24], [3, 3]); // 16 input channels, 24 output channels, 3x3 kernel size
+        let hidden_size = 24 * 22 * 22;
+        let fc1 = nn::LinearConfig::new(hidden_size, 32)
+            .with_bias(false)
+            .init();
+        let fc2 = nn::LinearConfig::new(32, NUM_CLASSES)
+            .with_bias(false)
+            .init();
+
+        let dropout = nn::DropoutConfig::new(0.5).init();
+
+        let instance = Self {
+            conv1: conv1,
+            conv2: conv2,
+            conv3: conv3,
+            dropout: dropout,
+            fc1: fc1,
+            fc2: fc2,
+            activation: nn::GELU::new(),
+        };
+        let record = BinBytesRecorder::<FullPrecisionSettings>::default()
+            .load(STATE_ENCODED.to_vec())
+            .expect("Failed to decode state");
+
+        instance.load_record(record)
+    }
+
+    /// Applies the forward pass of the neural network on the given input tensor.
+    ///
+    /// # Arguments
+    ///
+    /// * `input` - A 3-dimensional tensor of shape [batch_size, height, width].
+    ///
+    /// # Returns
+    ///
+    /// A 2-dimensional tensor of shape [batch_size, num_classes] containing the output of the neural network.
+    pub fn forward(&self, input: Tensor<B, 3>) -> Tensor<B, 2> {
+        // Get the dimensions of the input tensor
+        let [batch_size, height, width] = input.dims();
+        // Reshape the input tensor to have a shape of [batch_size, 1, height, width] and detach it
+        let x = input.reshape([batch_size, 1, height, width]).detach();
+        // Apply the first convolutional layer to the input tensor
+        let x = self.conv1.forward(x);
+        // Apply the second convolutional layer to the output of the first convolutional layer
+        let x = self.conv2.forward(x);
+        // Apply the third convolutional layer to the output of the second convolutional layer
+        let x = self.conv3.forward(x);
+
+        // Get the dimensions of the output tensor from the third convolutional layer
+        let [batch_size, channels, height, width] = x.dims();
+        // Reshape the output tensor to have a shape of [batch_size, channels*height*width]
+        let x = x.reshape([batch_size, channels * height * width]);
+
+        // Apply dropout to the output of the third convolutional layer
+        let x = self.dropout.forward(x);
+        // Apply the first fully connected layer to the output of the dropout layer
+        let x = self.fc1.forward(x);
+        // Apply the activation function to the output of the first fully connected layer
+        let x = self.activation.forward(x);
+
+        // Apply the second fully connected layer to the output of the activation function
+        self.fc2.forward(x)
+    }
+
+    pub fn inference(&mut self, input: &[f32]) -> Result<Vec<f32>, String> {
+        // Reshape from the 1D array to 3d tensor [batch, height, width]
+        let input: Tensor<B, 3> = Tensor::from_floats(input).reshape([1, 28, 28]);
+
+        // Normalize input: make between [0,1] and make the mean=0 and std=1
+        // values mean=0.1307, std=0.3081
+        // Source: https://github.com/pytorch/examples/blob/54f4572509891883a947411fd7239237dd2a39c3/mnist/main.py#L122
+        let input = ((input / 255) - 0.1307) / 0.3081;
+
+        // Run the tensor input through the model
+        let output: Tensor<B, 2> = self.forward(input);
+
+        // Convert the model output into probalibility distribution using softmax formula
+        let output = burn::tensor::activation::softmax(output, 1);
+
+        // Flatten oupuut tensor with [1,10] shape into boxed slice of [f32]
+        let output = output.into_data().convert::<f32>().value;
+
+        Ok(output)
+    }
+}
+
+/// A struct representing a convolutional block in a neural network model.
+#[derive(Module, Debug)]
+pub struct ConvBlock<B: Backend> {
+    /// A 2D convolutional layer.
+    conv: nn::conv::Conv2d<B>,
+    /// A batch normalization layer.
+    norm: BatchNorm<B, 2>,
+    /// A GELU activation function.
+    activation: nn::GELU,
+}
+
+/// A convolutional block with batch normalization and GELU activation.
+impl<B: Backend> ConvBlock<B> {
+    /// Creates a new `ConvBlock` with the given number of output channels and kernel size.
+    pub fn new(channels: [usize; 2], kernel_size: [usize; 2]) -> Self {
+        // Initialize a 2D convolutional layer with the given output channels and kernel size,
+        // and set the padding to "valid".
+        let conv = nn::conv::Conv2dConfig::new(channels, kernel_size)
+            .with_padding(PaddingConfig2d::Valid)
+            .init();
+
+        // Initialize a batch normalization layer with the number of channels in the second dimension of the output.
+        let norm = nn::BatchNormConfig::new(channels[1]).init();
+
+        // Create a new `ConvBlock` with the initialized convolutional and batch normalization layers,
+        // and a GELU activation function.
+        Self {
+            conv: conv,
+            norm: norm,
+            activation: nn::GELU::new(),
+        }
+    }
+
+    /// Applies the convolutional block to the given input tensor.
+    pub fn forward(&self, input: Tensor<B, 4>) -> Tensor<B, 4> {
+        // Apply the convolutional layer to the input tensor.
+        let x = self.conv.forward(input);
+
+        // Apply the batch normalization layer to the output of the convolutional layer.
+        let x = self.norm.forward(x);
+
+        // Apply the GELU activation function to the output of the batch normalization layer.
+        self.activation.forward(x)
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    #[cfg(feature = "ndarray")]
+    pub type Backend = burn::backend::NdArrayBackend<f32>;
+    #[test]
+    fn test_inference() {
+        let mut model = Model::<Backend>::new();
+        let output = model.inference(&[0.0; 28 * 28]).unwrap();
+        assert_eq!(output.len(), 10);
+    }
+}