DOCSEARCHbyAlgolia
Submit Docs

ort Rust Wrapper for ONNX Runtime

repository·main·Indexed 25 days ago

https://github.com/pykeio/ort

A safe Rust interface for hardware-accelerated inference and training on ONNX models. Primarily acting as a wrapper for Microsoft's ONNX Runtime 1.28, ort supports various execution providers and alternative pure-Rust backends such as ort-candle and ort-tract. It provides tools for on-device training via a Trainer API and supports deployment across datacenters, on-device execution, and WebAssembly (via ort-web).

Tokens
31.6K
Snippets
75
Records
154
Agent score
82%

What's inside ort

  1. Overview of `ort`

    main

    ort is a Rust interface for performing hardware-accelerated inference and training on machine learning models in the Open Neural Network Exchange (ONNX) format.

    It is primarily a wrapper for Microsoft's ONNX Runtime library, providing high performance and support for almost any hardware accelerator, but it also supports other pure-Rust runtimes. It is designed to be suitable for both datacenter deployment and on-device execution (e.g., for models exported from PyTorch, TensorFlow, Keras, scikit-learn, or PaddlePaddle).

  2. What is I/O Binding and when to use it

    main

    I/O Binding is an interface in ONNX Runtime that allows you to manually specify which inputs and outputs reside on which device and control when they are synchronized.

    When using non-CPU execution providers (like CUDA), copying data between the device and CPU can be a significant bottleneck. Use I/O binding to avoid these copies in scenarios where:

    1. An input does not change between runs (e.g., a style embedding).
    2. An output is used directly as an input to another model (or the same model) on the same device.

    By using I/O binding, you can keep data on the accelerator device, reducing expensive device-to-CPU transfers.

  3. Train models using the `ort` Trainer API

    main

    The ort library supports on-device training and fine-tuning without requiring an Execution Provider (EP).

    There are two ways to implement training:

    1. Manual Training Loop: Using the advanced Trainer API for full control over the training process (see train-clm example).
    2. Simple Training API: Using a Hugging Face-style interface that manages the loop for you (see train-clm-simple example).

    Example of the simple Trainer API usage:

    trainer.train(
    TrainingArguments::new(dataloader)
        .with_lr(7e-5)
        .with_max_steps(5000)
        .with_ckpt_strategy(CheckpointStrategy::Steps(500))
        .with_callbacks(LoggerCallback::new())
)?

  4. Synchronize data between Rust and ONNX Runtime WASM contexts

    main

    Because ort-web operates across two separate WebAssembly contexts, memory is not shared. You must manually synchronize data when moving it between the Rust context and the ONNX Runtime context.

    Input Tensors

    Do not use Tensor::new for inputs, as it allocates on the ONNX Runtime side and requires an unnecessary synchronization of empty data. Instead, use:

    • Tensor::from_array
    • TensorRef::from_array_view

    These methods create tensors that do not require synchronization.

    Output Tensors

    Session outputs are not synchronized automatically. To use output data in Rust, you must sync them. You can sync all outputs at once using ort_web::sync_outputs, or sync individual tensors using .sync(SyncDirection::Rust).await?.

    Sync Directions:

    • SyncDirection::Rust: Synchronizes data from the ONNX Runtime context to the Rust context. Use this after running a session to read outputs.
    • SyncDirection::Runtime: Synchronizes data from the Rust context to the ONNX Runtime context. Use this if you have modified a tensor in Rust and want the changes to be visible to the runtime.
    use ort_web::{TensorExt, SyncDirection};

// ... after session.run_async ...

let mut bounding_boxes = outputs.remove("bounding_boxes").unwrap();
bounding_boxes.sync(SyncDirection::Rust).await?;

// now we can use the data
let data = bounding_boxes.try_extract_tensor::<f32>()?;

  5. Understand prebuilt binary requirements and limitations

    main

    By default, ort downloads statically-linked ONNX Runtime binaries built by pyke.io to avoid long compilation times. Users should be aware of the following hardware and software requirements for these binaries:

    Hardware Requirements

    • x86-64 Architecture: All x86-64 binaries require the x86-64-v3 microarchitecture baseline. This includes:
      • Intel Haswell (Cores/Xeons after 2013)
      • Intel Gracemont (Atoms after Nov 2021)
      • AMD Excavator or later
      • Any AMD Ryzen processor
    • Note: Lower-cost Intel chips (like Pentiums) may not meet this requirement.

    Software Requirements

    • Linux: Binaries are compiled with Clang (not GCC) and depend on libc++.

    Execution Provider (EP) Combinations

    Prebuilt binaries are optimized for specific EP configurations. Some combinations are unavailable:

    • Windows: All builds include the DirectML EP.
    • macOS: All builds include the CoreML EP.
    • CUDA & TensorRT: These always ship together.
    • Conflict: You cannot enable both cuda and webgpu features simultaneously because no prebuilt binary exists containing both. Enabling both will result in a compile error.

  6. What are ONNX Values?

    main

    An ONNX value represents any data type that can be passed to or returned from an ONNX session or operator. There are three primary types of values:

    1. Tensors: Multi-dimensional arrays. This is the most common value type.
    2. Maps: Key-value pairs where keys map to a specific value type (similar to HashMap<K, V>).
    3. Sequences: Homogeneously-typed, dynamically-sized lists (similar to Vec<T>). Sequences can only contain tensors or maps of tensors.

  7. Use Views (TensorRef and TensorRefMut)

    main

    A view (or ref) is a borrowed variant of a value. They are used to pass data to a session without taking ownership or to reference external data.

    • TensorRef: A shared view of a tensor.
    • TensorRefMut: A mutable view of a tensor.

    Creating views of external data

    You can create views from existing ndarray arrays or raw slices. These views are bound to the lifetime of the original data.

    Passing values to a session

    You can pass values to session.run by value, by reference, or by view using the ort::inputs! macro.

    // Creating a view from an existing tensor
let my_tensor: ort::value::Tensor<f32> = Tensor::new(...)?;
let tensor_view: ort::value::TensorRef<'_, f32> = my_tensor.view();

// Creating a view from external data
let original_data = Array4::<f32>::from_shape_vec(...);
let tensor_view = TensorRef::from_array_view(original_data.view())?;

// Using views in session.run
let outputs = session.run(ort::inputs![
    "timestep" => timestep,
    "latents" => &latents,
    "text_embedding" => text_embedding.view()
])?;

  8. Use the `train-clm-simple` Trainer API

    main
    The train-clm-simple example demonstrates how to use ort's "simple" Trainer API. This API is designed to be high-level and similar to the Hugging Face Trainer or PyTorch Lightning APIs. Instead of implementing a manual training loop, you simply provide a data loader and training parameters, and the ort Trainer handles the training lifecycle for you.

  9. How Phi-3 Vision works with ONNX

    main

    Phi-3 Vision is a multimodal model that requires coordinating three interconnected ONNX models to process vision and language inputs. The workflow follows these steps:

    1. Image Processing: Preprocess the input image and pass it through the vision ONNX model to extract visual features.
    2. Text Embedding: Tokenize input text and process it with the text embedding ONNX model.
    3. Multimodal Fusion: Combine the visual features and text embeddings into a single input.
    4. Text Generation: Feed the combined input into the text generation ONNX model. The model generates text tokens autoregressively, using past key/value states to maintain context.

    Model configuration details are stored in data/genai_config.json.

  10. Enable Execution Providers and manage ONNX Runtime versions

    main

    To use specific hardware acceleration, you must enable the corresponding Cargo feature for that Execution Provider (EP).

    Additionally, the minimum ONNX Runtime version required by ort is controlled via api-* features. See the Multiversioning guide for details on managing version requirements.

  11. Handle Dynamic Values (DynValue)

    main

    Session outputs return DynValues, which are values whose exact type is not known at compile time. To use them, you have two main options:

    1. Try Extract: Use try_extract_* methods (e.g., try_extract_array). These return a Result and fail if the type is incompatible.
    2. Downcast: Convert the DynValue to a stronger type using .downcast(). This is preferred when you are certain of the type.

    Type Hierarchy

    • DynValue: Any type (Tensor, Map, or Sequence).
    • DynTensor / DynMap / DynSequence: The container type is known, but the element/key/value type is unknown.
    • Tensor<T> / Map<K, V> / Sequence<T>: Both container and element types are known.

  12. Quickstart: Load a model and perform inference

    main

    To run inference with ort, follow these steps:

    1. Convert your model: Ensure your model (from PyTorch, TensorFlow, etc.) is exported to the ONNX format.
    2. Load the model: Use Session::builder() to configure and load your .onnx file.
    3. Run inference: Use the run() method with the ort::inputs! macro to pass input tensors, then extract the results.

    Note: The example below uses GraphOptimizationLevel::Level3 for high performance and sets the number of intra-op threads.

    use ort::session::{builder::GraphOptimizationLevel, Session};

// 1. Load the model
let mut model = Session::builder()?
    .with_optimization_level(GraphOptimizationLevel::Level3)?
    .with_intra_threads(4)?
    .commit_from_file("yolov8m.onnx")?;

// 2. Perform inference
let outputs = model.run(ort::inputs!["image" => image])?;

// 3. Extract predictions
let predictions = outputs["output0"].try_extract_array::<f32>()?;