59 lines
2.0 KiB
Rust
59 lines
2.0 KiB
Rust
//! ONNX Runtime integration for RustyTorch++
|
|
//!
|
|
//! This crate provides high-performance ONNX model inference using ONNX Runtime,
|
|
//! with support for multiple execution providers (CPU, CUDA, CoreML, TensorRT).
|
|
//!
|
|
//! # Features
|
|
//!
|
|
//! - **CPU execution**: Always available, no additional dependencies
|
|
//! - **CUDA execution**: GPU acceleration on NVIDIA hardware (requires `cuda` feature)
|
|
//! - **CoreML execution**: Apple Neural Engine acceleration (requires `coreml` feature)
|
|
//! - **TensorRT execution**: Optimized inference on NVIDIA GPUs (requires `tensorrt` feature)
|
|
//!
|
|
//! # Example
|
|
//!
|
|
//! ```ignore
|
|
//! use rtx_onnx::{OnnxSession, OnnxSessionConfig};
|
|
//! use rtx_tensor::{Tensor, Device};
|
|
//! use std::collections::HashMap;
|
|
//!
|
|
//! // Load model with default config (auto-detect best execution provider)
|
|
//! let config = OnnxSessionConfig::default();
|
|
//! let session = OnnxSession::from_file("model.onnx", config)?;
|
|
//!
|
|
//! // Create input tensor
|
|
//! let input = Tensor::randn(vec![1, 3, 224, 224], &Device::CPU)?;
|
|
//!
|
|
//! // Run inference
|
|
//! let inputs = HashMap::from([("input".to_string(), &input)]);
|
|
//! let outputs = session.run(inputs)?;
|
|
//!
|
|
//! // Get output
|
|
//! let output = outputs.get("output").unwrap();
|
|
//! println!("Output shape: {:?}", output.shape());
|
|
//! ```
|
|
|
|
pub mod error;
|
|
pub mod execution_provider;
|
|
pub mod session;
|
|
pub mod tensor_bridge;
|
|
|
|
// Re-export main types
|
|
pub use error::{OnnxError, Result};
|
|
pub use execution_provider::{CpuOptions, ExecutionProviderType, detect_best_provider};
|
|
pub use session::{OnnxSession, OnnxSessionConfig, OptimizationLevel};
|
|
pub use tensor_bridge::{ort_to_rtx, rtx_to_ort};
|
|
|
|
// Re-export feature-gated execution provider options
|
|
#[cfg(feature = "cuda")]
|
|
pub use execution_provider::{ArenaExtendStrategy, CudaOptions, CudnnConvAlgoSearch};
|
|
|
|
#[cfg(feature = "coreml")]
|
|
pub use execution_provider::CoreMLOptions;
|
|
|
|
#[cfg(feature = "tensorrt")]
|
|
pub use execution_provider::TensorRTOptions;
|
|
|
|
#[cfg(feature = "directml")]
|
|
pub use execution_provider::DirectMLOptions;
|