From c8e9651907c3974eb558a4f7c1162d4308ac5c04 Mon Sep 17 00:00:00 2001 From: Ivan Smirnov Date: Sat, 28 Nov 2020 12:31:09 +0000 Subject: [PATCH 01/10] Add IDEA files to .gitignore --- .gitignore | 1 + 1 file changed, 1 insertion(+) diff --git a/.gitignore b/.gitignore index 76aa4374..b4bc6f68 100644 --- a/.gitignore +++ b/.gitignore @@ -5,3 +5,4 @@ onnxruntime.git Cargo.lock **/synset.txt +.idea/ From 10dedc7ee563a9ab7988e2e7a3c110dd37f0ed28 Mon Sep 17 00:00:00 2001 From: Ivan Smirnov Date: Sat, 28 Nov 2020 12:31:28 +0000 Subject: [PATCH 02/10] Reduce tracing level in tensor API to trace --- onnxruntime/src/tensor/ort_owned_tensor.rs | 6 +++--- onnxruntime/src/tensor/ort_tensor.rs | 6 +++--- 2 files changed, 6 insertions(+), 6 deletions(-) diff --git a/onnxruntime/src/tensor/ort_owned_tensor.rs b/onnxruntime/src/tensor/ort_owned_tensor.rs index 161fe105..fa46525a 100644 --- a/onnxruntime/src/tensor/ort_owned_tensor.rs +++ b/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(); diff --git a/onnxruntime/src/tensor/ort_tensor.rs b/onnxruntime/src/tensor/ort_tensor.rs index e4f3d93c..a8da24ab 100644 --- a/onnxruntime/src/tensor/ort_tensor.rs +++ b/onnxruntime/src/tensor/ort_tensor.rs @@ -3,7 +3,7 @@ use std::{fmt::Debug, ops::Deref}; use ndarray::Array; -use tracing::{debug, error}; +use tracing::{error, trace}; use onnxruntime_sys as sys; @@ -95,10 +95,10 @@ where T: TypeToTensorElementDataType + Debug + Clone, D: ndarray::Dimension, { - #[tracing::instrument] + #[tracing::instrument(level = "trace")] fn drop(&mut self) { // We need to let the C part free - debug!("Dropping Tensor."); + trace!("Dropping Tensor."); if self.c_ptr.is_null() { error!("Null pointer, not calling free."); } else { From 2f57def17a3af32e0658daa205287e3ab73a4de4 Mon Sep 17 00:00:00 2001 From: Ivan Smirnov Date: Sat, 28 Nov 2020 20:06:44 +0000 Subject: [PATCH 03/10] Implement DerefMut for OrtTensor --- onnxruntime/src/tensor/ort_tensor.rs | 12 +++++++++++- 1 file changed, 11 insertions(+), 1 deletion(-) diff --git a/onnxruntime/src/tensor/ort_tensor.rs b/onnxruntime/src/tensor/ort_tensor.rs index a8da24ab..7549b840 100644 --- a/onnxruntime/src/tensor/ort_tensor.rs +++ b/onnxruntime/src/tensor/ort_tensor.rs @@ -1,6 +1,6 @@ //! Module containing tensor with memory owned by Rust -use std::{fmt::Debug, ops::Deref}; +use std::{fmt::Debug, ops::Deref, ops::DerefMut}; use ndarray::Array; use tracing::{error, trace}; @@ -90,6 +90,16 @@ where } } +impl<'t, T, D> DerefMut for OrtTensor<'t, T, D> +where + T: TypeToTensorElementDataType + Debug + Clone, + D: ndarray::Dimension, +{ + fn deref_mut(&mut self) -> &mut ::Target { + &mut self.array + } +} + impl<'t, T, D> Drop for OrtTensor<'t, T, D> where T: TypeToTensorElementDataType + Debug + Clone, From fc3d1a12118e1dc68c2e853594024f0f36b658c3 Mon Sep 17 00:00:00 2001 From: Ivan Smirnov Date: Sat, 28 Nov 2020 20:07:09 +0000 Subject: [PATCH 04/10] Session::validate_input_shapes() doesn't need &mut --- onnxruntime/src/session.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/onnxruntime/src/session.rs b/onnxruntime/src/session.rs index 2cef12e5..adca2bb6 100644 --- a/onnxruntime/src/session.rs +++ b/onnxruntime/src/session.rs @@ -487,7 +487,7 @@ impl Session { // Tensor::from_array(self, array) // } - fn validate_input_shapes(&mut self, input_arrays: &[Array]) -> Result<()> + pub(crate) fn validate_input_shapes(&self, input_arrays: &[Array]) -> Result<()> where TIn: TypeToTensorElementDataType + Debug + Clone, D: ndarray::Dimension, From 6524450720a40150d28f350804b33eb8a1282990 Mon Sep 17 00:00:00 2001 From: Ivan Smirnov Date: Sat, 28 Nov 2020 20:07:23 +0000 Subject: [PATCH 05/10] Make session_ptr and memory_info pub(crate) --- onnxruntime/src/session.rs | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/onnxruntime/src/session.rs b/onnxruntime/src/session.rs index adca2bb6..df3896de 100644 --- a/onnxruntime/src/session.rs +++ b/onnxruntime/src/session.rs @@ -291,9 +291,9 @@ impl SessionBuilder { /// Type storing the session information, built from an [`Environment`](environment/struct.Environment.html) #[derive(Debug)] pub struct Session { - 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, /// Information about the ONNX's outputs as stored in loaded file From 4874fc180198f8fec65c6ebaee69f29b5399238b Mon Sep 17 00:00:00 2001 From: Ivan Smirnov Date: Sat, 28 Nov 2020 20:08:03 +0000 Subject: [PATCH 06/10] Add initial working implementation of Runner --- onnxruntime/src/error.rs | 4 + onnxruntime/src/lib.rs | 1 + onnxruntime/src/runner.rs | 231 ++++++++++++++++++++++++++++++++++++++ 3 files changed, 236 insertions(+) create mode 100644 onnxruntime/src/runner.rs diff --git a/onnxruntime/src/error.rs b/onnxruntime/src/error.rs index b63072ee..55020392 100644 --- a/onnxruntime/src/error.rs +++ b/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) diff --git a/onnxruntime/src/lib.rs b/onnxruntime/src/lib.rs index d326285d..2274fa3d 100644 --- a/onnxruntime/src/lib.rs +++ b/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; diff --git a/onnxruntime/src/runner.rs b/onnxruntime/src/runner.rs new file mode 100644 index 00000000..60104ba6 --- /dev/null +++ b/onnxruntime/src/runner.rs @@ -0,0 +1,231 @@ +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 Element for T {} + +fn names_to_ptrs(names: impl Iterator) -> Vec<*const i8> { + names + .map(|name| CString::new(name.clone()).unwrap().into_raw() as *const _) + .collect() +} + +fn compute_output_shapes( + input_arrays: &[Array], + outputs: &[Output], +) -> Vec> { + 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( + memory_info: &MemoryInfo, + arrays: impl IntoIterator>, +) -> Result>> { + Ok(arrays + .into_iter() + .map(|arr| OrtTensor::from_array(memory_info, arr)) + .collect::>>()?) +} + +fn tensors_to_ptr<'a, 's: 'a, T: Element, D: Dimension + 'a>( + tensors: impl IntoIterator>, +) -> 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>, +) -> Vec<*mut sys::OrtValue> { + tensors + .into_iter() + .map(|tensor| tensor.c_ptr as *mut _) + .collect() +} + +fn arrays_to_ort( + memory_info: &MemoryInfo, + arrays: impl IntoIterator>, +) -> Result<(Vec>, Vec<*const sys::OrtValue>)> { + let ort_tensors = arrays + .into_iter() + .map(|arr| OrtTensor::from_array(memory_info, arr)) + .collect::>>()?; + let ort_values = ort_tensors + .iter() + .map(|tensor| tensor.c_ptr as *const _) + .collect(); + Ok((ort_tensors, ort_values)) +} + +fn arrays_with_shapes(shapes: &[Vec]) -> Result>> { + Ok(shapes + .into_iter() + .map(|shape| Array::<_, IxDyn>::default(shape.clone()).into_dimensionality()) + .collect::>, _>>()?) +} + +pub struct Inputs<'r, 'a, T: Element, D: Dimension> { + tensors: &'a mut [OrtTensor<'r, T, D>], +} + +impl Inputs<'_, '_, T, D> {} + +impl Index for Inputs<'_, '_, T, D> { + type Output = Array; + + #[inline] + fn index(&self, index: usize) -> &Self::Output { + &(*self.tensors[index]) + } +} + +impl IndexMut 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 Outputs<'_, '_, T, D> {} + +impl Index for Outputs<'_, '_, T, D> { + type Output = Array; + + #[inline] + fn index(&self, index: usize) -> &Self::Output { + &(*self.tensors[index]) + } +} + +pub struct RunnerBuilder<'s, TIn: Element, DIn: Dimension> { + session: &'s Session, + input_arrays: Vec>, +} + +impl<'s, TIn: Element, DIn: Dimension> RunnerBuilder<'s, TIn, DIn> { + #[inline] + pub fn new( + session: &'s Session, + input_arrays: impl IntoIterator>, + ) -> Self { + Self { + session, + input_arrays: input_arrays.into_iter().collect(), + } + } + + #[inline] + pub fn with_output( + self, + ) -> Result> { + Runner::new(self.session, self.input_arrays) + } + + #[inline] + pub fn with_output_dyn(self) -> Result> { + Runner::new(self.session, self.input_arrays) + } +} + +pub struct Runner<'s, TIn: Element, DIn: Dimension, TOut: Element, DOut: Dimension> { + session: &'s Session, + input_names_ptr: Vec<*const i8>, + output_names_ptr: Vec<*const i8>, + input_ort_tensors: Vec>, + input_ort_values_ptr: Vec<*const sys::OrtValue>, + output_ort_tensors: Vec>, + output_ort_values_ptr: Vec<*mut sys::OrtValue>, +} + +impl<'s, TIn: Element, DIn: Dimension, TOut: Element, DOut: Dimension> + Runner<'s, TIn, DIn, TOut, DOut> +{ + pub fn new( + session: &'s Session, + input_arrays: impl IntoIterator>, + ) -> Result { + 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::>(); + 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)?) + } +} From a20bea6cf87977b21ce124645c72364ab563766b Mon Sep 17 00:00:00 2001 From: Ivan Smirnov Date: Sat, 28 Nov 2020 20:08:10 +0000 Subject: [PATCH 07/10] Add Session::make_runner() --- onnxruntime/src/session.rs | 10 +++++++++- 1 file changed, 9 insertions(+), 1 deletion(-) diff --git a/onnxruntime/src/session.rs b/onnxruntime/src/session.rs index df3896de..9c99837a 100644 --- a/onnxruntime/src/session.rs +++ b/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, @@ -361,6 +362,13 @@ impl Drop for Session { } impl Session { + pub fn make_runner>>( + &self, + input_arrays: I, + ) -> RunnerBuilder<'_, 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 From 6a1c591fa4d9ec7bf8de688301f3e6826c1f9e00 Mon Sep 17 00:00:00 2001 From: Ivan Smirnov Date: Sat, 28 Nov 2020 20:25:12 +0000 Subject: [PATCH 08/10] Properly drop name strings in Runner --- onnxruntime/src/runner.rs | 13 +++++++++++++ 1 file changed, 13 insertions(+) diff --git a/onnxruntime/src/runner.rs b/onnxruntime/src/runner.rs index 60104ba6..1755d2ea 100644 --- a/onnxruntime/src/runner.rs +++ b/onnxruntime/src/runner.rs @@ -229,3 +229,16 @@ impl<'s, TIn: Element, DIn: Dimension, TOut: Element, DOut: Dimension> .map_err(OrtError::Run)?) } } + +impl 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 _) }; + } + } +} From c7175272f3c1db35178786e0f0b8f69a7d2bbcac Mon Sep 17 00:00:00 2001 From: Nicolas Bigaouette Date: Thu, 7 Jan 2021 09:40:24 -0500 Subject: [PATCH 09/10] Fix lifetime errors following merge from master --- onnxruntime/src/runner.rs | 24 ++++++++++++------------ onnxruntime/src/session.rs | 2 +- 2 files changed, 13 insertions(+), 13 deletions(-) diff --git a/onnxruntime/src/runner.rs b/onnxruntime/src/runner.rs index 1755d2ea..fa46fd52 100644 --- a/onnxruntime/src/runner.rs +++ b/onnxruntime/src/runner.rs @@ -130,15 +130,15 @@ impl Index for Outputs<'_, '_, T, D> { } } -pub struct RunnerBuilder<'s, TIn: Element, DIn: Dimension> { - session: &'s Session, +pub struct RunnerBuilder<'s, 'a, TIn: Element, DIn: Dimension> { + session: &'s Session<'a>, input_arrays: Vec>, } -impl<'s, TIn: Element, DIn: Dimension> RunnerBuilder<'s, TIn, DIn> { +impl<'s, 'a, TIn: Element, DIn: Dimension> RunnerBuilder<'s, 'a, TIn, DIn> { #[inline] pub fn new( - session: &'s Session, + session: &'s Session<'a>, input_arrays: impl IntoIterator>, ) -> Self { Self { @@ -150,18 +150,18 @@ impl<'s, TIn: Element, DIn: Dimension> RunnerBuilder<'s, TIn, DIn> { #[inline] pub fn with_output( self, - ) -> Result> { + ) -> Result> { Runner::new(self.session, self.input_arrays) } #[inline] - pub fn with_output_dyn(self) -> Result> { + pub fn with_output_dyn(self) -> Result> { Runner::new(self.session, self.input_arrays) } } -pub struct Runner<'s, TIn: Element, DIn: Dimension, TOut: Element, DOut: Dimension> { - session: &'s Session, +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>, @@ -170,11 +170,11 @@ pub struct Runner<'s, TIn: Element, DIn: Dimension, TOut: Element, DOut: Dimensi output_ort_values_ptr: Vec<*mut sys::OrtValue>, } -impl<'s, TIn: Element, DIn: Dimension, TOut: Element, DOut: Dimension> - Runner<'s, TIn, DIn, TOut, DOut> +impl<'s, 'a, TIn: Element, DIn: Dimension, TOut: Element, DOut: Dimension> + Runner<'s, 'a, TIn, DIn, TOut, DOut> { pub fn new( - session: &'s Session, + session: &'s Session<'a>, input_arrays: impl IntoIterator>, ) -> Result { let input_names_ptr = names_to_ptrs(session.inputs.iter().map(|i| i.name.clone())); @@ -231,7 +231,7 @@ impl<'s, TIn: Element, DIn: Dimension, TOut: Element, DOut: Dimension> } impl Drop - for Runner<'_, TIn, DIn, TOut, DOut> + for Runner<'_, '_, TIn, DIn, TOut, DOut> { fn drop(&mut self) { for ptr in &self.input_names_ptr { diff --git a/onnxruntime/src/session.rs b/onnxruntime/src/session.rs index faa1c9ab..9bbfa16a 100644 --- a/onnxruntime/src/session.rs +++ b/onnxruntime/src/session.rs @@ -365,7 +365,7 @@ impl<'a> Session<'a> { pub fn make_runner>>( &self, input_arrays: I, - ) -> RunnerBuilder<'_, T, D> { + ) -> RunnerBuilder<'_, 'a, T, D> { RunnerBuilder::new(self, input_arrays) } From 6c0d254964696ed5877a4138bc6c42ef6519f24f Mon Sep 17 00:00:00 2001 From: Nicolas Bigaouette Date: Thu, 7 Jan 2021 09:54:46 -0500 Subject: [PATCH 10/10] Adapt sample example to runner API --- onnxruntime/examples/sample.rs | 36 +++++++++++++++++++++++++++------- 1 file changed, 29 insertions(+), 7 deletions(-) diff --git a/onnxruntime/examples/sample.rs b/onnxruntime/examples/sample.rs index d16d08da..1346fe6d 100644 --- a/onnxruntime/examples/sample.rs +++ b/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> = session.run(input_tensor_values)?; + // You can simply run the session with the input to get the output... + // let outputs: Vec> = 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::()?; + 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) { + let outputs: Outputs = runner.outputs(); + let output = &outputs[0]; + for i in 0..5 { + println!("Score for class [{}] = {}", i, output[[0, i, 0, 0]]); + } +}