From e355168338c0f96f2c71f7cc1bbb05e3fa09e99c Mon Sep 17 00:00:00 2001 From: Matt McCoy Date: Thu, 17 Sep 2026 20:40:00 -0400 Subject: [PATCH] Add ONNX config entries --- lib/ortex.ex | 12 +++++++++++- lib/ortex/model.ex | 22 ++++++++++++++++++++-- lib/ortex/native.ex | 2 +- native/ortex/src/lib.rs | 3 ++- native/ortex/src/model.rs | 18 +++++++++++++----- test/ortex_test.exs | 15 +++++++++++++++ 6 files changed, 62 insertions(+), 10 deletions(-) diff --git a/lib/ortex.ex b/lib/ortex.ex index 51b0e45..e242aff 100644 --- a/lib/ortex.ex +++ b/lib/ortex.ex @@ -18,14 +18,24 @@ defmodule Ortex do `config.exs` where `EXECUTION_PROVIDERS` is a list of strings of which execution providers to enable. + Session config entries can be passed as a fourth argument. These are handed + straight to ONNX Runtime as + [session configuration keys](https://github.com/microsoft/onnxruntime/blob/main/include/onnxruntime/core/session/onnxruntime_session_options_config_keys.h), + so any key ONNX Runtime understands works here. Values may be given as strings, + booleans (`true`/`false` become `"1"`/`"0"`), or numbers. + ## Examples iex> Ortex.load("./models/tinymodel.onnx") iex> Ortex.load("./models/tinymodel.onnx", [:cuda, :cpu]) iex> Ortex.load("./models/tinymodel.onnx", [:cpu], 0) + iex> Ortex.load("./models/tinymodel.onnx", [:cpu], 3, + ...> "session.intra_op.allow_spinning": false, + ...> "session.inter_op.allow_spinning": false + ...> ) """ - defdelegate load(path, eps \\ [:cpu], opt \\ 3), to: Ortex.Model + defdelegate load(path, eps \\ [:cpu], opt \\ 3, session_options \\ []), to: Ortex.Model @doc """ Run a forward pass through a model. diff --git a/lib/ortex/model.ex b/lib/ortex/model.ex index d8e550f..4bec0c8 100644 --- a/lib/ortex/model.ex +++ b/lib/ortex/model.ex @@ -22,8 +22,8 @@ defmodule Ortex.Model do defstruct [:reference] @doc false - def load(path, eps \\ [:cpu], opt \\ 3) do - case Ortex.Native.init(path, eps, opt) do + def load(path, eps \\ [:cpu], opt \\ 3, session_options \\ []) do + case Ortex.Native.init(path, eps, opt, normalize_session_options(session_options)) do {:error, msg} -> raise msg @@ -32,6 +32,24 @@ defmodule Ortex.Model do end end + # ONNX Runtime session config entries are string key/value pairs. Accept the + # natural Elixir spellings (atoms, booleans, numbers) and stringify them. + defp normalize_session_options(session_options) do + Enum.map(session_options, fn {key, value} -> + {config_key(key), config_value(value)} + end) + end + + defp config_key(key) when is_atom(key), do: Atom.to_string(key) + defp config_key(key) when is_binary(key), do: key + + defp config_value(true), do: "1" + defp config_value(false), do: "0" + defp config_value(value) when is_binary(value), do: value + defp config_value(value) when is_atom(value), do: Atom.to_string(value) + defp config_value(value) when is_integer(value), do: Integer.to_string(value) + defp config_value(value) when is_float(value), do: Float.to_string(value) + @doc false def run(%Ortex.Model{} = model, tensor) when not is_tuple(tensor) do run(model, {tensor}) diff --git a/lib/ortex/native.ex b/lib/ortex/native.ex index ee16c07..8b058a3 100644 --- a/lib/ortex/native.ex +++ b/lib/ortex/native.ex @@ -25,7 +25,7 @@ defmodule Ortex.Native do # When loading a NIF module, dummy clauses for all NIF function are required. # NIF dummies usually just error out when called when the NIF is not loaded, as that should never normally happen. - def init(_model_path, _execution_providers, _optimization_level), + def init(_model_path, _execution_providers, _optimization_level, _session_options), do: :erlang.nif_error(:nif_not_loaded) def run(_model, _inputs), do: :erlang.nif_error(:nif_not_loaded) diff --git a/native/ortex/src/lib.rs b/native/ortex/src/lib.rs index 0523105..656ce12 100644 --- a/native/ortex/src/lib.rs +++ b/native/ortex/src/lib.rs @@ -22,9 +22,10 @@ fn init( model_path: String, eps: Vec, opt: i32, + session_options: Vec<(String, String)>, ) -> NifResult> { let eps = utils::map_eps(env, eps); - let model = model::init(model_path, eps, opt) + let model = model::init(model_path, eps, opt, session_options) .map_err(|e| rustler::Error::Term(Box::new(e.to_string())))?; Ok(ResourceArc::new(model)) } diff --git a/native/ortex/src/model.rs b/native/ortex/src/model.rs index f5e04b3..845f2ce 100644 --- a/native/ortex/src/model.rs +++ b/native/ortex/src/model.rs @@ -4,7 +4,7 @@ //! # Examples //! //! ``` -//! let model = init("./models/resnet50.onnx", vec![])?; +//! let model = init("./models/resnet50.onnx", vec![], 3, vec![])?; //! let (inputs, outputs) = show(model)?; //! ``` @@ -30,19 +30,27 @@ pub struct OrtexModel { unsafe impl Sync for OrtexModel {} /// Creates a model given the path to the model and vector of execution providers. -/// The execution providers are Atoms from Erlang/Elixir. +/// The execution providers are Atoms from Erlang/Elixir. `session_options` are +/// key/value pairs passed straight through to ONNX Runtime as session config +/// entries, e.g. `("session.intra_op.allow_spinning", "0")`. pub fn init( model_path: String, eps: Vec, opt: i32, + session_options: Vec<(String, String)>, ) -> Result { // TODO: send tracing logs to erlang/elixir _somehow_ // tracing_subscriber::fmt::init(); - let session = Session::builder()? + let mut builder = Session::builder()? .with_optimization_level(map_opt_level(opt))? - .with_execution_providers(eps)? - .commit_from_file(model_path)?; + .with_execution_providers(eps)?; + + for (key, value) in session_options.iter() { + builder = builder.with_config_entry(key, value)?; + } + + let session = builder.commit_from_file(model_path)?; let state = OrtexModel { session }; Ok(state) diff --git a/test/ortex_test.exs b/test/ortex_test.exs index bb41f8b..25bb061 100644 --- a/test/ortex_test.exs +++ b/test/ortex_test.exs @@ -23,6 +23,21 @@ defmodule OrtexTest do assert result |> Nx.backend_transfer() |> Nx.argmax(axis: 1) == Nx.tensor([499]) end + test "session config entries" do + model = + Ortex.load("./models/tinymodel.onnx", [:cpu], 3, + "session.intra_op.allow_spinning": false, + "session.inter_op.allow_spinning": 0, + "session.disable_prepacking": "1" + ) + + {%Nx.Tensor{shape: {1, 10}}, %Nx.Tensor{shape: {1, 10}}, %Nx.Tensor{shape: {1, 10}}} = + Ortex.run(model, { + Nx.broadcast(0, {1, 100}) |> Nx.as_type(:s32), + Nx.broadcast(0.0, {1, 100}) |> Nx.as_type(:f32) + }) + end + test "Nx.Serving with tinymodel" do model = Ortex.load("./models/tinymodel.onnx")