diff --git a/build.rs b/build.rs index dcd59f0..8637156 100644 --- a/build.rs +++ b/build.rs @@ -28,8 +28,6 @@ fn main() { .parse_callbacks(Box::new(bindgen::CargoCallbacks)) .whitelist_function("OrtGetApiBase") .whitelist_function("OrtSessionOptionsAppendExecutionProvider_CPU") - .whitelist_function("OrtSessionOptionsAppendExecutionProvider_OpenVINO") - .whitelist_function("OrtSessionOptionsAppendExecutionProvider_CUDA") .whitelist_function("OrtSessionOptionsAppendExecutionProvider_Tensorrt") .whitelist_var("ORT_.*") .whitelist_recursively(true) diff --git a/cbits/ort.h b/cbits/ort.h index 7e79865..52d81b3 100644 --- a/cbits/ort.h +++ b/cbits/ort.h @@ -1,6 +1,4 @@ #include ORT_API_STATUS(OrtSessionOptionsAppendExecutionProvider_CPU, _In_ OrtSessionOptions* options, int use_arena) ORT_ALL_ARGS_NONNULL; -ORT_API_STATUS(OrtSessionOptionsAppendExecutionProvider_CUDA, _In_ OrtSessionOptions* options, int device_id); ORT_API_STATUS(OrtSessionOptionsAppendExecutionProvider_Tensorrt, _In_ OrtSessionOptions* options, int device_id); -ORT_API_STATUS(OrtSessionOptionsAppendExecutionProvider_OpenVINO, _In_ OrtSessionOptions* options, const char* device_id); diff --git a/examples/yolo.rs b/examples/yolo.rs index 0bae7d9..7a45130 100644 --- a/examples/yolo.rs +++ b/examples/yolo.rs @@ -30,147 +30,6 @@ struct Opt { runs: usize, } -// #[derive(Debug, Clone)] -// struct ModelInfo { -// pub path: PathBuf, -// pub name: String, -// pub description: String, -// pub version: i64, -// pub input_shape: Vec, -// pub output_shape: Vec, -// pub input_name: String, -// pub output_name: String, -// } - -// struct OnnxModelConfig { -// pub model_path: PathBuf, -// pub parallel: bool, -// pub use_cuda: bool, -// pub use_cpu: bool, -// pub use_tensorrt: bool, -// pub cpu_use_arena:bool, -// pub tensorrt_device_index: i32, -// pub cuda_device_index: i32, -// pub logging_level: LoggingLevel, -// pub logging_id: String, -// } - -// struct OnnxModel { -// info: Arc, -// env: Env, -// options: SessionOptions, -// session: Session, -// } - -// impl OnnxModel { -// pub fn new(config: OnnxModelConfig) -> Result { -// let env = Env::new(config.logging_level, &config.logging_id)?; -// let mut options = SessionOptions::new()?; - -// if config.parallel { -// options.set_execution_mode(ExecutionMode::Parallel)?; -// } - -// if config.use_cpu { -// options.add_cpu(config.cpu_use_arena); -// } - -// if config.use_cuda { -// options.add_cuda(config.cuda_device_index); -// } - -// if config.use_tensorrt { -// options.add_tensorrt(config.tensorrt_device_index); -// } - -// let session = Session::new(&env, config.model_path.to_str().unwrap(), &options)?; -// let metadata = session.metadata(); -// let inputs = session.inputs(); -// let outputs = session.outputs(); - -// if inputs.len() != 1 { -// return Err(Error::ModelWrongInputsCount); -// } - -// if outputs.len() != 1 { -// return Err(Error::ModelWrongOutputsCount); -// } - -// let input = session.input(0); -// let input_info = input -// .tensor_info() -// .or_else(|| Error::ModelWrongInputType)?; -// let input_name = input.name().to_string(); -// let input_shape: Vec = input_info -// .symbolic_dims() -// .map(|d| match d { -// SymbolicDim::Symbolic(_) => -1i64, -// SymbolicDim::Fixed(x) => x as i64, -// }) -// .collect(); - -// let output = session.output(0); -// let output_info = output -// .tensor_info() -// .or_else(|| Error::ModelWrongInputType)?; -// let output_name = output.name().to_string(); -// let output_shape: Vec = output_info -// .symbolic_dims() -// .map(|d| match d { -// SymbolicDim::Symbolic(_) => -1i64, -// SymbolicDim::Fixed(x) => x as i64, -// }) -// .collect(); - -// let info = ModelInfo { -// path: path.into(), -// name: format!("{} exported by {}", metadata.graph_name(), metadata.producer_name()), -// description: metadata.description().to_string(), -// version: metadata.version(), -// input_name, -// input_shape, -// output_name, -// output_shape, -// }; - -// Ok(OnnxModel { -// info: Arc::new(info), -// env, -// options, -// session, -// }) -// } -// } - - -// pub enum ModelCommand { -// AddOnnxModel(usize, OnnxModelConfig), -// Inference(usize, ) -// } - -// struct ModelPool { -// counter: usize, -// sender: Sender, -// } - -// impl ModelPool { -// pub fn add(&mut self, name: &str, file: Path) -> Result<&AsyncModel, Error> { - -// } - -// pub fn get_model(&self, name: &str) -> Option<&AsyncModel> { - -// } -// } - - -// struct AsyncModel { - -// } - - - - use std::collections::HashMap; fn key_val_parse(str: &str) -> HashMap { @@ -251,7 +110,7 @@ fn tensor_with_size( println!("{:?} {} {:?}", ty, name, dims); match ty { Float => match name { - "input" => Box::new(Tensor::::new(&dims, load_image("/home/andrey/Images/me.jpg", dims[2], dims[3])).unwrap()), + "input" => Box::new(Tensor::::new(&dims, load_image("/data/andrey_/Images/me.jpg", dims[2], dims[3])).unwrap()), _ => Box::new(Tensor::::init(&dims, 0.0).unwrap()), }, Int64 => Box::new(Tensor::::init(&dims, 0).unwrap()), @@ -268,7 +127,7 @@ fn main() -> Result<()> { // so.set_execution_mode(ExecutionMode::Parallel)?; // so.add_tensorrt(0); - // so.add_cuda(0); + so.add_cuda(0); // so.add_cpu(true); let mut map = if let Some(dims) = &opt.dims { diff --git a/src/lib.rs b/src/lib.rs index b1dfd30..a917470 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -8,7 +8,7 @@ pub mod sys; // Re-export enums pub use sys::{ AllocatorType, ErrorCode, ExecutionMode, GraphOptimizationLevel, LoggingLevel, MemType, - OnnxTensorElementDataType, OnnxType, + OnnxTensorElementDataType, OnnxType, CUDAProviderOptions, CudnnConvAlgoSearch, OpenVINOProviderOptions }; #[macro_use] @@ -173,19 +173,29 @@ impl SessionOptions { } pub fn add_cuda(&self, device_id: i32) { - - let so = self.raw; - let status = unsafe { - crate::sys::SessionOptionsAppendExecutionProvider_CUDA(so, device_id) + let cuda_options = CUDAProviderOptions { + device_id, + cudnn_conv_algo_search: CudnnConvAlgoSearch::Default, + cuda_mem_limit: u64::MAX, + arena_extend_strategy: 0, + do_copy_in_default_stream: 1, }; - if !status.is_null() { - panic!("!!!"); - } + call!(@unsafe @expect SessionOptionsAppendExecutionProvider_CUDA, self.raw, &cuda_options); + } + + pub fn add_open_vino(&self, device_type: &CStr, device_id: &CStr) { + let openvino_options = OpenVINOProviderOptions { + device_type: device_type.as_ptr(), + enable_vpu_fast_compile: 0, + device_id: device_id.as_ptr(), + num_of_threads: 0, + }; + + call!(@unsafe @expect SessionOptionsAppendExecutionProvider_OpenVINO, self.raw, &openvino_options); } pub fn add_tensorrt(&self, device_id: i32) { - let so = self.raw; let status = unsafe { crate::sys::SessionOptionsAppendExecutionProvider_Tensorrt(so, device_id)