ML_RUNNER

A small library that aims to provide a fast and device agnostic way of running python ML models in embedded platforms using rust

Features

• A python library to parse ML models into an exported json format. Exports from :
  • ONNX models
• A rust library to run the exported models on any platforms, with optimized functions depending on the platforms toolset

The library supports complex models with multiple inputs and outputs, and with multiple paths inside the model.

Supported layers in python library

Supported activation functions

Supported layers in Rust library

Supported activation functions

Usage

Python library

To export a simple example onnx model, you can run the following :

poetry run python main.py --model-path simple_linear_model.onnx --output-path export.json

or from your own code

from ml_runner_exporter.onnx_exporter import export_onnx

output_model = export_onnx(model_path)

You can then save the exported model as a json file

Rust library

To run a model using the rust library, first export it to json using the python library as described in the python library section.

You can then import it and run it using :

let model = match Model::from_json(&model_json) {
    Ok(m) => m,
    Err(e) => {
        eprintln!("Error parsing model JSON: {}", e);
        return;
    }
};

// Shape the input to what the model requires
let input = vec![...];
println!("Input: {:?}", input);

// Run forward pass
match model.forward(&input) {
    Ok(output) => println!("Output: {:?}", output),
    Err(e) => eprintln!("Error during forward pass: {}", e),
}

See the main.rs file for a concrete example.

To run the library, you can use one of the following targets :

• cargo run to run the default target (no optimization)
• cargo run --features simd to run with SIMD optimizations (requires wide crate)
• cargo run --features blas to run with BLAS optimizations (requires specific software depending on your OS)
• To compile in release mode, use cargo run -r (add the feature you want after)

Testing

To run the tests, you can use cargo test --all-features to run all tests

Benchmarks

See benchmark.md

Roadmap

These are the next features that I would like to implement, in no specific order

• Add activation functions support in simd
• Add support for more layer types
• Add support to other optimization backends (cuda, ...)
• Export both the python library and rust library to pip/crates.io for easier usage
• Add support for non float models (right now only float is supported, we should allow export and run of models in int)
• Implement parser for pytorch and tensorflow models
  • I will focus mainly on tensorflow support at first, since pytorch has a good onnx exporter