diff --git a/.changeset/wakeword_enable_graph_optimization.md b/.changeset/wakeword_enable_graph_optimization.md new file mode 100644 index 000000000..feb4716ac --- /dev/null +++ b/.changeset/wakeword_enable_graph_optimization.md @@ -0,0 +1,12 @@ +--- +livekit-wakeword: patch +--- + +# Enable graph optimization for the tract backend + +Sessions were built with `Session::builder()` without setting a graph optimization +level. The `ort-tract` backend used on every target except aarch64 Windows runs +tract's `into_optimized()` only when the session requests an optimization level, so +wake word inference ran the unoptimized graph. Requesting `Level3` — ONNX Runtime's +own default, so the native backend is unaffected — made `predict()` over a 2 s window +7.4x faster (534.5 ms to 72.5 ms median, Apple M-series release build). diff --git a/livekit-wakeword/src/lib.rs b/livekit-wakeword/src/lib.rs index 619869de9..dc0ef3cd4 100644 --- a/livekit-wakeword/src/lib.rs +++ b/livekit-wakeword/src/lib.rs @@ -16,7 +16,10 @@ use std::path::Path; #[cfg(use_tract)] use std::sync::Once; -use ort::session::Session; +use ort::session::{ + builder::{GraphOptimizationLevel, SessionBuilder}, + Session, +}; #[cfg(use_tract)] static INIT_TRACT: Once = Once::new(); @@ -74,15 +77,25 @@ pub(crate) fn to_resampler_rate(hz: u32) -> Result Result { + Ok(Session::builder()?.with_optimization_level(GraphOptimizationLevel::Level3)?) +} + pub(crate) fn build_session_from_memory(bytes: &[u8]) -> Result { #[cfg(use_tract)] ensure_tract_backend(); - Ok(Session::builder()?.commit_from_memory(bytes)?) + Ok(session_builder()?.commit_from_memory(bytes)?) } pub(crate) fn build_session_from_file(path: impl AsRef) -> Result { #[cfg(use_tract)] ensure_tract_backend(); let bytes = std::fs::read(path)?; - Ok(Session::builder()?.commit_from_memory(&bytes)?) + Ok(session_builder()?.commit_from_memory(&bytes)?) }