Runner

.gitignore

@@ -5,3 +5,4 @@
 onnxruntime.git
 Cargo.lock
 **/synset.txt
+.idea/

onnxruntime/examples/sample.rs

@@ -1,8 +1,11 @@
 #![forbid(unsafe_code)]
 
 use onnxruntime::{
- environment::Environment, ndarray::Array, tensor::OrtOwnedTensor, GraphOptimizationLevel,
- LoggingLevel,
+ environment::Environment,
+ ndarray::{Array, Ix4, IxDyn},
+ runner::{Outputs, Runner},
+ tensor::OrtOwnedTensor,
+ GraphOptimizationLevel, LoggingLevel,
 };
 use tracing::Level;
 use tracing_subscriber::FmtSubscriber;
@@ -61,12 +64,31 @@ fn run() -> Result<(), Error> {
  .unwrap();
  let input_tensor_values = vec![array];
 
- let outputs: Vec<OrtOwnedTensor<f32, _>> = session.run(input_tensor_values)?;
+ // You can simply run the session with the input to get the output...
+ // let outputs: Vec<OrtOwnedTensor<f32, _>> = session.run(input_tensor_values)?;
 
- assert_eq!(outputs[0].shape(), output0_shape.as_slice());
- for i in 0..5 {
- println!("Score for class [{}] = {}", i, outputs[0][[0, i, 0, 0]]);
- }
+ // Or, you can build a runner to pre-allocate the output
+ let mut runner = session
+ .make_runner(input_tensor_values)
+ .with_output::<f32, Ix4>()?;
+ runner.execute()?;
+
+ print_runner_outputs(&runner);
+
+ // Since the runner now owns the input and keep it alive, we can access it
+ // and modify it without reallocation.
+ *(&mut runner.inputs()[0]) *= 2.0f32;
+ runner.execute()?;
+
+ print_runner_outputs(&runner);
 
  Ok(())
 }
+
+fn print_runner_outputs(runner: &Runner<f32, IxDyn, f32, Ix4>) {
+ let outputs: Outputs<f32, Ix4> = runner.outputs();
+ let output = &outputs[0];
+ for i in 0..5 {
+ println!("Score for class [{}] = {}", i, output[[0, i, 0, 0]]);
+ }
+}

onnxruntime/src/error.rs

@@ -4,6 +4,7 @@ use std::{io, path::PathBuf};
 
 use thiserror::Error;
 
+use ndarray::ShapeError;
 use onnxruntime_sys as sys;
 
 use crate::{char_p_to_string, g_ort};
@@ -91,6 +92,9 @@ pub enum OrtError {
  /// Attempt to build a Rust `CString` from a null pointer
  #[error("Failed to build CString when original contains null: {0}")]
  CStringNulError(#[from] std::ffi::NulError),
+ /// Output dimensionality mismatch
+ #[error("Output dimensionality mismatch: {0}")]
+ OutputDimensionalityMismatch(#[from] ShapeError),
 }
 
 /// Error used when dimensions of input (from model and from inference call)

onnxruntime/src/lib.rs

@@ -125,6 +125,7 @@ pub mod download;
 pub mod environment;
 pub mod error;
 mod memory;
+pub mod runner;
 pub mod session;
 pub mod tensor;
 

onnxruntime/src/runner.rs

@@ -0,0 +1,244 @@
+use std::ffi::CString;
+use std::fmt::Debug;
+use std::ops::{Index, IndexMut};
+
+use ndarray::{Array, Dimension, IxDyn};
+
+use onnxruntime_sys as sys;
+
+use crate::error::{status_to_result, OrtError};
+use crate::memory::MemoryInfo;
+use crate::session::{Output, Session};
+use crate::tensor::OrtTensor;
+use crate::{g_ort, Result, TypeToTensorElementDataType};
+
+pub trait Element: 'static + Clone + Debug + TypeToTensorElementDataType + Default {}
+
+impl<T: 'static + Clone + Debug + TypeToTensorElementDataType + Default> Element for T {}
+
+fn names_to_ptrs(names: impl Iterator<Item = String>) -> Vec<*const i8> {
+ names
+ .map(|name| CString::new(name.clone()).unwrap().into_raw() as *const _)
+ .collect()
+}
+
+fn compute_output_shapes<TIn, DIn: Dimension>(
+ input_arrays: &[Array<TIn, DIn>],
+ outputs: &[Output],
+) -> Vec<Vec<usize>> {
+ outputs
+ .iter()
+ .enumerate()
+ .map(|(idx, output)| {
+ output
+ .dimensions
+ .iter()
+ .enumerate()
+ .map(|(jdx, dim)| match dim {
+ None => input_arrays[idx].shape()[jdx],
+ Some(d) => *d as usize,
+ })
+ .collect()
+ })
+ .collect()
+}
+
+fn arrays_to_tensors<T: Element, D: Dimension>(
+ memory_info: &MemoryInfo,
+ arrays: impl IntoIterator<Item = Array<T, D>>,
+) -> Result<Vec<OrtTensor<T, D>>> {
+ Ok(arrays
+ .into_iter()
+ .map(|arr| OrtTensor::from_array(memory_info, arr))
+ .collect::<Result<Vec<_>>>()?)
+}
+
+fn tensors_to_ptr<'a, 's: 'a, T: Element, D: Dimension + 'a>(
+ tensors: impl IntoIterator<Item = &'a OrtTensor<'s, T, D>>,
+) -> Vec<*const sys::OrtValue> {
+ tensors
+ .into_iter()
+ .map(|tensor| tensor.c_ptr as *const _)
+ .collect()
+}
+
+fn tensors_to_mut_ptr<'a, 's: 'a, T: Element, D: Dimension + 'a>(
+ tensors: impl IntoIterator<Item = &'a mut OrtTensor<'s, T, D>>,
+) -> Vec<*mut sys::OrtValue> {
+ tensors
+ .into_iter()
+ .map(|tensor| tensor.c_ptr as *mut _)
+ .collect()
+}
+
+fn arrays_to_ort<T: Element, D: Dimension>(
+ memory_info: &MemoryInfo,
+ arrays: impl IntoIterator<Item = Array<T, D>>,
+) -> Result<(Vec<OrtTensor<T, D>>, Vec<*const sys::OrtValue>)> {
+ let ort_tensors = arrays
+ .into_iter()
+ .map(|arr| OrtTensor::from_array(memory_info, arr))
+ .collect::<Result<Vec<_>>>()?;
+ let ort_values = ort_tensors
+ .iter()
+ .map(|tensor| tensor.c_ptr as *const _)
+ .collect();
+ Ok((ort_tensors, ort_values))
+}
+
+fn arrays_with_shapes<T: Element, D: Dimension>(shapes: &[Vec<usize>]) -> Result<Vec<Array<T, D>>> {
+ Ok(shapes
+ .into_iter()
+ .map(|shape| Array::<_, IxDyn>::default(shape.clone()).into_dimensionality())
+ .collect::<std::result::Result<Vec<Array<T, D>>, _>>()?)
+}
+
+pub struct Inputs<'r, 'a, T: Element, D: Dimension> {
+ tensors: &'a mut [OrtTensor<'r, T, D>],
+}
+
+impl<T: Element, D: Dimension> Inputs<'_, '_, T, D> {}
+
+impl<T: Element, D: Dimension> Index<usize> for Inputs<'_, '_, T, D> {
+ type Output = Array<T, D>;
+
+ #[inline]
+ fn index(&self, index: usize) -> &Self::Output {
+ &(*self.tensors[index])
+ }
+}
+
+impl<T: Element, D: Dimension> IndexMut<usize> for Inputs<'_, '_, T, D> {
+ #[inline]
+ fn index_mut(&mut self, index: usize) -> &mut Self::Output {
+ &mut (*self.tensors[index])
+ }
+}
+
+pub struct Outputs<'r, 'a, T: Element, D: Dimension> {
+ tensors: &'a [OrtTensor<'r, T, D>],
+}
+
+impl<T: Element, D: Dimension> Outputs<'_, '_, T, D> {}
+
+impl<T: Element, D: Dimension> Index<usize> for Outputs<'_, '_, T, D> {
+ type Output = Array<T, D>;
+
+ #[inline]
+ fn index(&self, index: usize) -> &Self::Output {
+ &(*self.tensors[index])
+ }
+}
+
+pub struct RunnerBuilder<'s, 'a, TIn: Element, DIn: Dimension> {
+ session: &'s Session<'a>,
+ input_arrays: Vec<Array<TIn, DIn>>,
+}
+
+impl<'s, 'a, TIn: Element, DIn: Dimension> RunnerBuilder<'s, 'a, TIn, DIn> {
+ #[inline]
+ pub fn new(
+ session: &'s Session<'a>,
+ input_arrays: impl IntoIterator<Item = Array<TIn, DIn>>,
+ ) -> Self {
+ Self {
+ session,
+ input_arrays: input_arrays.into_iter().collect(),
+ }
+ }
+
+ #[inline]
+ pub fn with_output<TOut: Element, DOut: Dimension>(
+ self,
+ ) -> Result<Runner<'s, 'a, TIn, DIn, TOut, DOut>> {
+ Runner::new(self.session, self.input_arrays)
+ }
+
+ #[inline]
+ pub fn with_output_dyn<TOut: Element>(self) -> Result<Runner<'s, 'a, TIn, DIn, TOut, IxDyn>> {
+ Runner::new(self.session, self.input_arrays)
+ }
+}
+
+pub struct Runner<'s, 'a, TIn: Element, DIn: Dimension, TOut: Element, DOut: Dimension> {
+ session: &'s Session<'a>,
+ input_names_ptr: Vec<*const i8>,
+ output_names_ptr: Vec<*const i8>,
+ input_ort_tensors: Vec<OrtTensor<'s, TIn, DIn>>,
+ input_ort_values_ptr: Vec<*const sys::OrtValue>,
+ output_ort_tensors: Vec<OrtTensor<'s, TOut, DOut>>,
+ output_ort_values_ptr: Vec<*mut sys::OrtValue>,
+}
+
+impl<'s, 'a, TIn: Element, DIn: Dimension, TOut: Element, DOut: Dimension>
+ Runner<'s, 'a, TIn, DIn, TOut, DOut>
+{
+ pub fn new(
+ session: &'s Session<'a>,
+ input_arrays: impl IntoIterator<Item = Array<TIn, DIn>>,
+ ) -> Result<Self> {
+ let input_names_ptr = names_to_ptrs(session.inputs.iter().map(|i| i.name.clone()));
+ let output_names_ptr = names_to_ptrs(session.outputs.iter().map(|o| o.name.clone()));
+ let input_arrays = input_arrays.into_iter().collect::<Vec<_>>();
+ session.validate_input_shapes(&input_arrays)?;
+ let output_shapes = compute_output_shapes(&input_arrays, &session.outputs);
+ let output_arrays = arrays_with_shapes::<_, DOut>(&output_shapes)?;
+ let input_ort_tensors = arrays_to_tensors(&session.memory_info, input_arrays)?;
+ let input_ort_values_ptr = tensors_to_ptr(&input_ort_tensors);
+ let mut output_ort_tensors = arrays_to_tensors(&session.memory_info, output_arrays)?;
+ let output_ort_values_ptr = tensors_to_mut_ptr(&mut output_ort_tensors);
+ Ok(Self {
+ session,
+ input_names_ptr,
+ output_names_ptr,
+ input_ort_tensors,
+ input_ort_values_ptr,
+ output_ort_tensors,
+ output_ort_values_ptr,
+ })
+ }
+
+ #[inline]
+ pub fn inputs(&mut self) -> Inputs<'s, '_, TIn, DIn> {
+ Inputs {
+ tensors: self.input_ort_tensors.as_mut_slice(),
+ }
+ }
+
+ #[inline]
+ pub fn outputs(&'s self) -> Outputs<'s, '_, TOut, DOut> {
+ Outputs {
+ tensors: self.output_ort_tensors.as_slice(),
+ }
+ }
+
+ #[inline]
+ pub fn execute(&mut self) -> Result<()> {
+ Ok(status_to_result(unsafe {
+ g_ort().Run.unwrap()(
+ self.session.session_ptr,
+ std::ptr::null() as _,
+ self.input_names_ptr.as_ptr(),
+ self.input_ort_values_ptr.as_ptr(),
+ self.input_ort_values_ptr.len() as _,
+ self.output_names_ptr.as_ptr(),
+ self.output_names_ptr.len() as _,
+ self.output_ort_values_ptr.as_mut_ptr(),
+ )
+ })
+ .map_err(OrtError::Run)?)
+ }
+}
+
+impl<TIn: Element, DIn: Dimension, TOut: Element, DOut: Dimension> Drop
+ for Runner<'_, '_, TIn, DIn, TOut, DOut>
+{
+ fn drop(&mut self) {
+ for ptr in &self.input_names_ptr {
+ let _s = unsafe { CString::from_raw(*ptr as _) };
+ }
+ for ptr in &self.output_names_ptr {
+ let _s = unsafe { CString::from_raw(*ptr as _) };
+ }
+ }
+}

onnxruntime/src/session.rs

@@ -10,7 +10,7 @@ use std::os::windows::ffi::OsStrExt;
 #[cfg(feature = "model-fetching")]
 use std::env;
 
-use ndarray::Array;
+use ndarray::{Array, Dimension};
 use tracing::{debug, error};
 
 use onnxruntime_sys as sys;
@@ -21,6 +21,7 @@ use crate::{
  error::{status_to_result, NonMatchingDimensionsError, OrtError, Result},
  g_ort,
  memory::MemoryInfo,
+ runner::{Element, RunnerBuilder},
  tensor::{
  ort_owned_tensor::{OrtOwnedTensor, OrtOwnedTensorExtractor},
  OrtTensor,
@@ -291,9 +292,9 @@ impl<'a> SessionBuilder<'a> {
 #[derive(Debug)]
 pub struct Session<'a> {
  env: &'a Environment,
- session_ptr: *mut sys::OrtSession,
+ pub(crate) session_ptr: *mut sys::OrtSession,
  allocator_ptr: *mut sys::OrtAllocator,
- memory_info: MemoryInfo,
+ pub(crate) memory_info: MemoryInfo,
  /// Information about the ONNX's inputs as stored in loaded file
  pub inputs: Vec<Input>,
  /// Information about the ONNX's outputs as stored in loaded file
@@ -361,6 +362,13 @@ impl<'a> Drop for Session<'a> {
 }
 
 impl<'a> Session<'a> {
+ pub fn make_runner<T: Element, D: Dimension, I: IntoIterator<Item = Array<T, D>>>(
+ &self,
+ input_arrays: I,
+ ) -> RunnerBuilder<'_, 'a, T, D> {
+ RunnerBuilder::new(self, input_arrays)
+ }
+
  /// Run the input data through the ONNX graph, performing inference.
  ///
  /// Note that ONNX models can have multiple inputs; a `Vec<_>` is thus
@@ -487,7 +495,7 @@ impl<'a> Session<'a> {
  // Tensor::from_array(self, array)
  // }
 
- fn validate_input_shapes<TIn, D>(&mut self, input_arrays: &[Array<TIn, D>]) -> Result<()>
+ pub(crate) fn validate_input_shapes<TIn, D>(&self, input_arrays: &[Array<TIn, D>]) -> Result<()>
  where
  TIn: TypeToTensorElementDataType + Debug + Clone,
  D: ndarray::Dimension,

onnxruntime/src/tensor/ort_owned_tensor.rs

@@ -3,7 +3,7 @@
 use std::{fmt::Debug, ops::Deref};
 
 use ndarray::{Array, ArrayView};
-use tracing::debug;
+use tracing::trace;
 
 use onnxruntime_sys as sys;
 
@@ -124,9 +124,9 @@ where
  D: ndarray::Dimension,
  'm: 't, // 'm outlives 't
 {
- #[tracing::instrument]
+ #[tracing::instrument(level = "trace")]
  fn drop(&mut self) {
- debug!("Dropping OrtOwnedTensor.");
+ trace!("Dropping OrtOwnedTensor.");
  unsafe { g_ort().ReleaseValue.unwrap()(self.tensor_ptr) }
 
  self.tensor_ptr = std::ptr::null_mut();