livekit-wakeword: runtime-configurable ONNX session options

.changeset/wakeword_runtime_session_options.md

@@ -0,0 +1,20 @@
+---
+livekit-wakeword: minor
+---
+
+# Runtime-configurable ONNX session options
+
+`WakeWordModel::with_session_options` takes a `SessionOptions` describing how the
+crate creates its ONNX sessions — graph optimization level, intra/inter-op threads,
+sequential execution, thread spinning, and arbitrary session config entries —
+matching the `sess_options` parameter the Python SDK takes. It applies to the two
+bundled feature extraction models and to every wake word classifier, including ones
+a later `load_model` call adds.
+
+`WakeWordModel::new` keeps its signature and uses `SessionOptions::default()`, whose
+`GraphOptimizationLevel::Level3` matches ONNX Runtime's own default. That default
+matters: the `ort-tract` backend used on every target except aarch64 Windows runs
+tract's `into_optimized()` only when a level is requested, so requesting one made
+`predict()` over a 2 s window 6.2x faster (329 ms to 53 ms median, Apple M-series
+release build). tract implements no other session option, so the remaining fields
+are skipped there rather than failing session creation.

livekit-wakeword/src/embedding.rs

@@ -16,7 +16,7 @@ use ndarray::{Array, Array1};
 use ort::session::Session;
 use ort::value::Tensor;
 
-use crate::{build_session_from_memory, WakeWordError};
+use crate::{build_session_from_memory, SessionOptions, WakeWordError};
 
 const MODEL_BYTES: &[u8] = include_bytes!("../onnx/embedding_model.onnx");
 
@@ -33,8 +33,8 @@ pub struct EmbeddingModel {
 }
 
 impl EmbeddingModel {
-    pub fn new() -> Result<Self, WakeWordError> {
-        Ok(Self { session: build_session_from_memory(MODEL_BYTES)? })
+    pub fn new(options: &SessionOptions) -> Result<Self, WakeWordError> {
+        Ok(Self { session: build_session_from_memory(MODEL_BYTES, options)? })
     }
 
     // Run the embedding model on mel spectrogram features and return the embedding.

livekit-wakeword/src/lib.rs

@@ -16,7 +16,7 @@ use std::path::Path;
 #[cfg(use_tract)]
 use std::sync::Once;
 
-use ort::session::Session;
+use ort::session::{builder::SessionBuilder, Session};
 
 #[cfg(use_tract)]
 static INIT_TRACT: Once = Once::new();
@@ -32,6 +32,8 @@ pub(crate) mod embedding;
 pub(crate) mod melspectrogram;
 pub mod wakeword;
 
+/// Graph optimizations requested when an ONNX session is created.
+pub use ort::session::builder::GraphOptimizationLevel;
 pub use wakeword::WakeWordModel;
 
 #[derive(Debug, thiserror::Error)]
@@ -74,15 +76,121 @@ pub(crate) fn to_resampler_rate(hz: u32) -> Result<resampler::SampleRate, WakeWo
     }
 }
 
-pub(crate) fn build_session_from_memory(bytes: &[u8]) -> Result<Session, WakeWordError> {
+/// How every ONNX session a [`WakeWordModel`] creates is configured: the two
+/// bundled feature extraction models and each wake word classifier. Mirrors the
+/// `sess_options` parameter the Python SDK takes.
+///
+/// [`Default`] is what [`WakeWordModel::new`] uses. Pass a value of your own to
+/// [`WakeWordModel::with_session_options`] to trade latency for CPU, which is
+/// useful when detection runs as a background process on a small machine:
+///
+/// ```no_run
+/// # use livekit_wakeword::{SessionOptions, WakeWordModel};
+/// # fn main() -> Result<(), livekit_wakeword::WakeWordError> {
+/// let options = SessionOptions {
+///     intra_threads: Some(1),
+///     inter_threads: Some(1),
+///     parallel_execution: Some(false),
+///     intra_op_spinning: Some(false),
+///     inter_op_spinning: Some(false),
+///     ..Default::default()
+/// };
+/// let model = WakeWordModel::with_session_options(&["hey_livekit.onnx"], 16000, options)?;
+/// # Ok(())
+/// # }
+/// ```
+///
+/// `optimization_level` is the only option the `ort-tract` backend used on every
+/// target except aarch64 Windows implements; the rest are skipped there rather than
+/// failing session creation, and since tract runs single-threaded without
+/// spin-waiting, they already describe what it does.
+#[derive(Clone, Debug)]
+pub struct SessionOptions {
+    /// Graph optimizations to apply when the session is created. Matters more than
+    /// it looks: `ort-tract` runs tract's `into_optimized()` only when some level is
+    /// requested, and wake word inference measured several times slower without it.
+    pub optimization_level: GraphOptimizationLevel,
+    /// Threads used to run a single operator. `Some(1)` keeps each session on the
+    /// calling thread.
+    pub intra_threads: Option<usize>,
+    /// Threads used to run operators in parallel, when `parallel_execution` is on.
+    pub inter_threads: Option<usize>,
+    /// Whether independent operators may run in parallel. `Some(false)` is ONNX
+    /// Runtime's sequential execution mode.
+    pub parallel_execution: Option<bool>,
+    /// Whether intra-op threads may spin before blocking. `Some(false)` trades a
+    /// little latency for markedly less CPU while idle.
+    pub intra_op_spinning: Option<bool>,
+    /// Whether inter-op threads may spin before blocking.
+    pub inter_op_spinning: Option<bool>,
+    /// Session config entries applied verbatim, for anything the fields above do
+    /// not cover.
+    pub config_entries: Vec<(String, String)>,
+}
+
+impl Default for SessionOptions {
+    /// Requests [`GraphOptimizationLevel::Level3`], ONNX Runtime's own default, and
+    /// leaves every other option to the backend.
+    fn default() -> Self {
+        Self {
+            optimization_level: GraphOptimizationLevel::Level3,
+            intra_threads: None,
+            inter_threads: None,
+            parallel_execution: None,
+            intra_op_spinning: None,
+            inter_op_spinning: None,
+            config_entries: Vec::new(),
+        }
+    }
+}
+
+impl SessionOptions {
+    // The optimization level is the only session option `ort-tract` implements;
+    // every other one resolves to `ort-sys`' stub API and reports `NotImplemented`.
+    // Since tract runs single-threaded without spin-waiting, skipping those is
+    // closer to what the caller asked for than refusing to build the session.
     #[cfg(use_tract)]
-    ensure_tract_backend();
-    Ok(Session::builder()?.commit_from_memory(bytes)?)
+    fn apply(&self, builder: SessionBuilder) -> Result<SessionBuilder, WakeWordError> {
+        Ok(builder.with_optimization_level(self.optimization_level)?)
+    }
+
+    #[cfg(not(use_tract))]
+    fn apply(&self, builder: SessionBuilder) -> Result<SessionBuilder, WakeWordError> {
+        let mut builder = builder.with_optimization_level(self.optimization_level)?;
+        if let Some(threads) = self.intra_threads {
+            builder = builder.with_intra_threads(threads)?;
+        }
+        if let Some(threads) = self.inter_threads {
+            builder = builder.with_inter_threads(threads)?;
+        }
+        if let Some(parallel) = self.parallel_execution {
+            builder = builder.with_parallel_execution(parallel)?;
+        }
+        if let Some(enable) = self.intra_op_spinning {
+            builder = builder.with_intra_op_spinning(enable)?;
+        }
+        if let Some(enable) = self.inter_op_spinning {
+            builder = builder.with_inter_op_spinning(enable)?;
+        }
+        for (key, value) in &self.config_entries {
+            builder = builder.with_config_entry(key, value)?;
+        }
+        Ok(builder)
+    }
 }
 
-pub(crate) fn build_session_from_file(path: impl AsRef<Path>) -> Result<Session, WakeWordError> {
+pub(crate) fn build_session_from_memory(
+    bytes: &[u8],
+    options: &SessionOptions,
+) -> Result<Session, WakeWordError> {
     #[cfg(use_tract)]
     ensure_tract_backend();
-    let bytes = std::fs::read(path)?;
-    Ok(Session::builder()?.commit_from_memory(&bytes)?)
+    Ok(options.apply(Session::builder()?)?.commit_from_memory(bytes)?)
+}
+
+pub(crate) fn build_session_from_file(
+    path: impl AsRef<Path>,
+    options: &SessionOptions,
+) -> Result<Session, WakeWordError> {
+    build_session_from_memory(&std::fs::read(path)?, options)
 }

livekit-wakeword/src/melspectrogram.rs

@@ -16,7 +16,7 @@ use ndarray::{Array1, Array2, Axis};
 use ort::session::Session;
 use ort::value::Tensor;
 
-use crate::{build_session_from_memory, WakeWordError};
+use crate::{build_session_from_memory, SessionOptions, WakeWordError};
 
 const MODEL_BYTES: &[u8] = include_bytes!("../onnx/melspectrogram.onnx");
 
@@ -32,8 +32,8 @@ pub struct MelspectrogramModel {
 }
 
 impl MelspectrogramModel {
-    pub fn new() -> Result<Self, WakeWordError> {
-        Ok(Self { session: build_session_from_memory(MODEL_BYTES)? })
+    pub fn new(options: &SessionOptions) -> Result<Self, WakeWordError> {
+        Ok(Self { session: build_session_from_memory(MODEL_BYTES, options)? })
     }
 
     // Run the melspectrogram model on normalized f32 audio and return mel features.

livekit-wakeword/src/wakeword.rs

@@ -23,8 +23,8 @@ use resampler::{Attenuation, Latency, ResamplerFir, SampleRate};
 use crate::embedding::EmbeddingModel;
 use crate::melspectrogram::MelspectrogramModel;
 use crate::{
-    build_session_from_file, to_resampler_rate, WakeWordError, EMBEDDING_STRIDE, EMBEDDING_WINDOW,
-    MIN_EMBEDDINGS,
+    build_session_from_file, to_resampler_rate, SessionOptions, WakeWordError, EMBEDDING_STRIDE,
+    EMBEDDING_WINDOW, MIN_EMBEDDINGS,
 };
 
 struct Resampler {
@@ -45,6 +45,7 @@ pub struct WakeWordModel {
     emb_model: EmbeddingModel,
     classifiers: HashMap<String, Session>,
     resampler: Option<Resampler>,
+    session_options: SessionOptions,
 }
 
 impl WakeWordModel {
@@ -53,7 +54,23 @@ impl WakeWordModel {
     /// The recommended sample rate is 16 kHz. Other supported rates
     /// (22050, 32000, 44100, 48000, 88200, 96000, 176400, 192000, 384000 Hz)
     /// are resampled internally to 16 kHz.
+    ///
+    /// ONNX sessions are created with [`SessionOptions::default`]; use
+    /// [`with_session_options`](Self::with_session_options) to tune them.
     pub fn new(models: &[impl AsRef<Path>], sample_rate: u32) -> Result<Self, WakeWordError> {
+        Self::with_session_options(models, sample_rate, SessionOptions::default())
+    }
+
+    /// Create a new wake word model, tuning the ONNX sessions it creates.
+    ///
+    /// `session_options` applies to the bundled feature extraction models, to every
+    /// classifier in `models`, and to any classifier a later
+    /// [`load_model`](Self::load_model) call adds.
+    pub fn with_session_options(
+        models: &[impl AsRef<Path>],
+        sample_rate: u32,
+        session_options: SessionOptions,
+    ) -> Result<Self, WakeWordError> {
         let resampler = if sample_rate != 16000 {
             let input_rate = to_resampler_rate(sample_rate)?;
             // FIR resampler: 64-sample latency (~1.3ms at 48kHz) with 90dB
@@ -72,10 +89,11 @@ impl WakeWordModel {
         };
 
         let mut wakeword = Self {
-            mel_model: MelspectrogramModel::new()?,
-            emb_model: EmbeddingModel::new()?,
+            mel_model: MelspectrogramModel::new(&session_options)?,
+            emb_model: EmbeddingModel::new(&session_options)?,
             classifiers: HashMap::new(),
             resampler,
+            session_options,
         };
 
         for path in models {
@@ -87,7 +105,8 @@ impl WakeWordModel {
 
     /// Load a wake word classifier ONNX model from disk.
     ///
-    /// If `model_name` is `None`, the file stem is used as the classifier name.
+    /// If `model_name` is `None`, the file stem is used as the classifier name. The
+    /// session is created with the [`SessionOptions`] this model was built with.
     pub fn load_model(
         &mut self,
         model_path: impl AsRef<Path>,
@@ -103,7 +122,7 @@ impl WakeWordModel {
             None => path.file_stem().and_then(|s| s.to_str()).unwrap_or("unknown").to_string(),
         };
 
-        let session = build_session_from_file(path)?;
+        let session = build_session_from_file(path, &self.session_options)?;
         self.classifiers.insert(name, session);
         Ok(())
     }

livekit-wakeword/tests/integration.rs

@@ -14,7 +14,7 @@
 
 use std::path::PathBuf;
 
-use livekit_wakeword::{WakeWordModel, SAMPLE_RATE};
+use livekit_wakeword::{GraphOptimizationLevel, SessionOptions, WakeWordModel, SAMPLE_RATE};
 
 mod common;
 
@@ -82,3 +82,65 @@ fn test_negative_wav_below_threshold() {
         "expected negative sample score ({score}) < threshold ({THRESHOLD})"
     );
 }
+
+/// Session options supplied at runtime must not change what the model predicts,
+/// including the options a backend cannot honour (`ort-tract` implements only the
+/// graph optimization level and skips the rest).
+#[test]
+fn test_session_options_preserve_scores() {
+    let (sample_rate, samples) = common::read_wav("positive.wav");
+
+    let options = SessionOptions {
+        intra_threads: Some(1),
+        inter_threads: Some(1),
+        parallel_execution: Some(false),
+        intra_op_spinning: Some(false),
+        inter_op_spinning: Some(false),
+        config_entries: vec![("session.use_env_allocators".to_string(), "1".to_string())],
+        ..Default::default()
+    };
+
+    let mut default = WakeWordModel::new(&[classifier_path()], sample_rate).unwrap();
+    let mut tuned =
+        WakeWordModel::with_session_options(&[classifier_path()], sample_rate, options).unwrap();
+
+    assert_eq!(
+        default.predict(&samples).unwrap()["hey_livekit"],
+        tuned.predict(&samples).unwrap()["hey_livekit"]
+    );
+}
+
+/// Every optimization level builds a model that still detects the wake word.
+#[test]
+fn test_optimization_levels() {
+    let (sample_rate, samples) = common::read_wav("positive.wav");
+
+    for level in [
+        GraphOptimizationLevel::Disable,
+        GraphOptimizationLevel::Level1,
+        GraphOptimizationLevel::Level3,
+    ] {
+        let options = SessionOptions { optimization_level: level, ..Default::default() };
+        let mut model =
+            WakeWordModel::with_session_options(&[classifier_path()], sample_rate, options)
+                .unwrap();
+        let score = model.predict(&samples).unwrap()["hey_livekit"];
+        assert!(score > THRESHOLD, "{level:?} scored {score}");
+    }
+}
+
+/// A classifier loaded after construction inherits the model's session options.
+#[test]
+fn test_load_model_inherits_session_options() {
+    let (sample_rate, samples) = common::read_wav("positive.wav");
+
+    let no_models: &[PathBuf] = &[];
+    let options = SessionOptions {
+        optimization_level: GraphOptimizationLevel::Disable,
+        ..Default::default()
+    };
+    let mut model = WakeWordModel::with_session_options(no_models, sample_rate, options).unwrap();
+    model.load_model(classifier_path(), Some("late")).unwrap();
+
+    assert!(model.predict(&samples).unwrap()["late"] > THRESHOLD);
+}