diff --git a/.changeset/wakeword_runtime_session_options.md b/.changeset/wakeword_runtime_session_options.md new file mode 100644 index 000000000..71cce0e0b --- /dev/null +++ b/.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. diff --git a/livekit-wakeword/src/embedding.rs b/livekit-wakeword/src/embedding.rs index 350a5dd4d..a707409ee 100644 --- a/livekit-wakeword/src/embedding.rs +++ b/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 { - Ok(Self { session: build_session_from_memory(MODEL_BYTES)? }) + pub fn new(options: &SessionOptions) -> Result { + Ok(Self { session: build_session_from_memory(MODEL_BYTES, options)? }) } // Run the embedding model on mel spectrogram features and return the embedding. diff --git a/livekit-wakeword/src/lib.rs b/livekit-wakeword/src/lib.rs index 619869de9..cf22edd62 100644 --- a/livekit-wakeword/src/lib.rs +++ b/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 Result { +/// 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, + /// Threads used to run operators in parallel, when `parallel_execution` is on. + pub inter_threads: Option, + /// Whether independent operators may run in parallel. `Some(false)` is ONNX + /// Runtime's sequential execution mode. + pub parallel_execution: Option, + /// 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, + /// Whether inter-op threads may spin before blocking. + pub inter_op_spinning: Option, + /// 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 { + Ok(builder.with_optimization_level(self.optimization_level)?) + } + + #[cfg(not(use_tract))] + fn apply(&self, builder: SessionBuilder) -> Result { + 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) -> Result { +pub(crate) fn build_session_from_memory( + bytes: &[u8], + options: &SessionOptions, +) -> Result { #[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, + options: &SessionOptions, +) -> Result { + build_session_from_memory(&std::fs::read(path)?, options) } diff --git a/livekit-wakeword/src/melspectrogram.rs b/livekit-wakeword/src/melspectrogram.rs index e677513a6..f0718dd37 100644 --- a/livekit-wakeword/src/melspectrogram.rs +++ b/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 { - Ok(Self { session: build_session_from_memory(MODEL_BYTES)? }) + pub fn new(options: &SessionOptions) -> Result { + Ok(Self { session: build_session_from_memory(MODEL_BYTES, options)? }) } // Run the melspectrogram model on normalized f32 audio and return mel features. diff --git a/livekit-wakeword/src/wakeword.rs b/livekit-wakeword/src/wakeword.rs index 1002be6e1..bc658162a 100644 --- a/livekit-wakeword/src/wakeword.rs +++ b/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, resampler: Option, + 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], sample_rate: u32) -> Result { + 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], + sample_rate: u32, + session_options: SessionOptions, + ) -> Result { 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, @@ -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(()) } diff --git a/livekit-wakeword/tests/integration.rs b/livekit-wakeword/tests/integration.rs index 9afb3271d..2fcbd656e 100644 --- a/livekit-wakeword/tests/integration.rs +++ b/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); +}