andrey/onnxruntime-rs
1
0
Fork 0
Code Issues 1 Pull Requests Projects Releases Wiki Activity

Fix bindings for CUDA Provider

Browse Source
This commit is contained in:
Andrey Tkachenko 2021-01-22 13:45:06 +04:00
parent 25785a1ada
commit 8041b81588
4 changed files with 21 additions and 156 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
build.rs
View File

@ -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)

2
cbits/ort.h
View File

@ -1,6 +1,4 @@
#include <onnxruntime/core/session/onnxruntime_c_api.h>
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);

145
examples/yolo.rs
View File

@ -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<i64>,
// pub output_shape: Vec<i64>,
// 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<ModelInfo>,
// env: Env,
// options: SessionOptions,
// session: Session,
// }
// impl OnnxModel {
// pub fn new(config: OnnxModelConfig) -> Result<Self, Error> {
// 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<i64> = 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<i64> = 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<ModelCommand>,
// }
// 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<String, usize> {
@ -251,7 +110,7 @@ fn tensor_with_size(
println!("{:?} {} {:?}", ty, name, dims);
match ty {
Float => match name {
"input" => Box::new(Tensor::<f32>::new(&dims, load_image("/home/andrey/Images/me.jpg", dims[2], dims[3])).unwrap()),
"input" => Box::new(Tensor::<f32>::new(&dims, load_image("/data/andrey_/Images/me.jpg", dims[2], dims[3])).unwrap()),
_ => Box::new(Tensor::<f32>::init(&dims, 0.0).unwrap()),
},
Int64 => Box::new(Tensor::<i64>::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 {

26
src/lib.rs
View File

@ -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)