From 5534bb688e443a387619c9cb63605853de74ba91 Mon Sep 17 00:00:00 2001 From: Peter Nehrer Date: Thu, 16 May 2024 09:45:18 -0400 Subject: [PATCH 1/5] Expose setting input/getting output tensors by index --- crates/openvino/src/request.rs | 29 ++++++++++++++++++++++++++--- 1 file changed, 26 insertions(+), 3 deletions(-) diff --git a/crates/openvino/src/request.rs b/crates/openvino/src/request.rs index 2ca3d9a..67b2ec3 100644 --- a/crates/openvino/src/request.rs +++ b/crates/openvino/src/request.rs @@ -1,9 +1,10 @@ use crate::tensor::Tensor; use crate::{cstr, drop_using_function, try_unsafe, util::Result}; use openvino_sys::{ - ov_infer_request_free, ov_infer_request_get_tensor, ov_infer_request_infer, - ov_infer_request_set_tensor, ov_infer_request_start_async, ov_infer_request_t, - ov_infer_request_wait_for, + ov_infer_request_free, ov_infer_request_get_output_tensor_by_index, + ov_infer_request_get_tensor, ov_infer_request_infer, + ov_infer_request_set_input_tensor_by_index, ov_infer_request_set_tensor, + ov_infer_request_start_async, ov_infer_request_t, ov_infer_request_wait_for, }; /// See [`InferRequest`](https://docs.openvino.ai/2023.3/api/c_cpp_api/group__ov__infer__request__c__api.html). @@ -20,6 +21,7 @@ impl InferRequest { pub(crate) fn from_ptr(ptr: *mut ov_infer_request_t) -> Self { Self { ptr } } + /// Assign a [`Tensor`] to the input on the model. pub fn set_tensor(&mut self, name: &str, tensor: &Tensor) -> Result<()> { try_unsafe!(ov_infer_request_set_tensor( @@ -41,6 +43,27 @@ impl InferRequest { Ok(Tensor::from_ptr(tensor)) } + /// Assing an input [`Tensor`] to the model by its index. + pub fn set_input_tensor_by_index(&mut self, index: usize, tensor: &Tensor) -> Result<()> { + try_unsafe!(ov_infer_request_set_input_tensor_by_index( + self.ptr, + index, + tensor.as_ptr() + ))?; + Ok(()) + } + + /// Retrieve an output [`Tensor`] from the model by its index. + pub fn get_output_tensor_by_index(&self, index: usize) -> Result { + let mut tensor = std::ptr::null_mut(); + try_unsafe!(ov_infer_request_get_output_tensor_by_index( + self.ptr, + index, + std::ptr::addr_of_mut!(tensor) + ))?; + Ok(Tensor::from_ptr(tensor)) + } + /// Execute the inference request. pub fn infer(&mut self) -> Result<()> { try_unsafe!(ov_infer_request_infer(self.ptr)) From df4e3080d700fc180734bd8671e7745758ef4345 Mon Sep 17 00:00:00 2001 From: Peter Nehrer Date: Thu, 16 May 2024 13:18:33 -0400 Subject: [PATCH 2/5] Support loading non-string models from memory --- crates/openvino/src/core.rs | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/crates/openvino/src/core.rs b/crates/openvino/src/core.rs index aadfc93..a2ee614 100644 --- a/crates/openvino/src/core.rs +++ b/crates/openvino/src/core.rs @@ -10,6 +10,8 @@ use openvino_sys::{ ov_core_read_model, ov_core_read_model_from_memory_buffer, ov_core_t, }; +use std::os::raw::c_char; + /// See [`Core`](https://docs.openvino.ai/2023.3/api/c_cpp_api/group__ov__core__c__api.html). pub struct Core { ptr: *mut ov_core_t, @@ -53,15 +55,15 @@ impl Core { /// Read model with model and weights loaded in memory. pub fn read_model_from_buffer( &mut self, - model_str: &str, - weights_buffer: &Tensor, + model_str: &[u8], + weights_buffer: Option<&Tensor>, ) -> Result { let mut ptr = std::ptr::null_mut(); try_unsafe!(ov_core_read_model_from_memory_buffer( self.ptr, - cstr!(model_str), + model_str.as_ptr().cast::(), model_str.len(), - weights_buffer.as_ptr(), + weights_buffer.map_or(std::ptr::null(), |tensor| tensor.as_ptr().cast_const()), std::ptr::addr_of_mut!(ptr) ))?; Ok(Model::from_ptr(ptr)) From c6ff11a82aa6e41b9bf73592a114ce17ce2a7f52 Mon Sep 17 00:00:00 2001 From: Peter Nehrer Date: Thu, 16 May 2024 14:30:14 -0400 Subject: [PATCH 3/5] Improve ergonomics --- crates/openvino/src/core.rs | 17 ++++++++++++++--- 1 file changed, 14 insertions(+), 3 deletions(-) diff --git a/crates/openvino/src/core.rs b/crates/openvino/src/core.rs index a2ee614..afb2c0c 100644 --- a/crates/openvino/src/core.rs +++ b/crates/openvino/src/core.rs @@ -53,17 +53,19 @@ impl Core { } /// Read model with model and weights loaded in memory. - pub fn read_model_from_buffer( + pub fn read_model_from_buffer<'a, T: Into>>( &mut self, model_str: &[u8], - weights_buffer: Option<&Tensor>, + weights_buffer: T, ) -> Result { let mut ptr = std::ptr::null_mut(); try_unsafe!(ov_core_read_model_from_memory_buffer( self.ptr, model_str.as_ptr().cast::(), model_str.len(), - weights_buffer.map_or(std::ptr::null(), |tensor| tensor.as_ptr().cast_const()), + weights_buffer + .into() + .map_or(std::ptr::null(), |tensor| tensor.as_ptr().cast_const()), std::ptr::addr_of_mut!(ptr) ))?; Ok(Model::from_ptr(ptr)) @@ -87,9 +89,18 @@ impl Core { #[cfg(test)] mod core_tests { use super::*; + #[test] fn test_new() { let core = Core::new(); assert!(core.is_ok()); } + + #[test] + fn test_load_onnx_from_buffer() { + let model = b"\x08\x07\x12\nonnx-wally:j\n*\n\x06inputs\x12\x07outputs\x1a\ridentity_node\"\x08Identity\x12\x0bno-op-modelZ\x16\n\x06inputs\x12\x0c\n\n\x08\x01\x12\x06\n\x00\n\x02\x08\x02b\x17\n\x07outputs\x12\x0c\n\n\x08\x01\x12\x06\n\x00\n\x02\x08\x02B\x02\x10\x0c"; + let mut core = Core::new().unwrap(); + let model = core.read_model_from_buffer(model, None); + assert!(model.is_ok()); + } } From 095c739567549394fa5bd415743bf7809db89959 Mon Sep 17 00:00:00 2001 From: Peter Nehrer Date: Fri, 17 May 2024 16:12:49 -0400 Subject: [PATCH 4/5] Use explicit Option arg --- crates/openvino/src/core.rs | 8 +++----- 1 file changed, 3 insertions(+), 5 deletions(-) diff --git a/crates/openvino/src/core.rs b/crates/openvino/src/core.rs index afb2c0c..39008cd 100644 --- a/crates/openvino/src/core.rs +++ b/crates/openvino/src/core.rs @@ -53,19 +53,17 @@ impl Core { } /// Read model with model and weights loaded in memory. - pub fn read_model_from_buffer<'a, T: Into>>( + pub fn read_model_from_buffer( &mut self, model_str: &[u8], - weights_buffer: T, + weights_buffer: Option<&Tensor>, ) -> Result { let mut ptr = std::ptr::null_mut(); try_unsafe!(ov_core_read_model_from_memory_buffer( self.ptr, model_str.as_ptr().cast::(), model_str.len(), - weights_buffer - .into() - .map_or(std::ptr::null(), |tensor| tensor.as_ptr().cast_const()), + weights_buffer.map_or(std::ptr::null(), |tensor| tensor.as_ptr().cast_const()), std::ptr::addr_of_mut!(ptr) ))?; Ok(Model::from_ptr(ptr)) From 437f24130f027ffcadeca4cecdb5a0fb0fe8d633 Mon Sep 17 00:00:00 2001 From: Peter Nehrer Date: Fri, 17 May 2024 17:26:27 -0400 Subject: [PATCH 5/5] Remove 2022.3.0 from matrix --- .github/workflows/main.yml | 6 +----- 1 file changed, 1 insertion(+), 5 deletions(-) diff --git a/.github/workflows/main.yml b/.github/workflows/main.yml index ccbd5d4..6a21fbc 100644 --- a/.github/workflows/main.yml +++ b/.github/workflows/main.yml @@ -25,12 +25,8 @@ jobs: apt: [false] # We also spot-check that things work when installing from APT by adding to the matrix: see # https://docs.github.com/en/actions/using-jobs/using-a-matrix-for-your-jobs#expanding-or-adding-matrix-configurations - # APT install and check oldest supported version 2023.2.0 - include: - - os: ubuntu-22.04 - version: 2022.3.0 - apt: false # APT install and check latest supported version 2024.1.0 + include: - os: ubuntu-22.04 version: 2024.1.0 apt: true