83 lines
2.2 KiB
Rust
83 lines
2.2 KiB
Rust
//! Burn ML Framework Integration for `RustyTorch`++
|
|
//!
|
|
//! This crate provides integration with the [Burn](https://github.com/tracel-ai/burn)
|
|
//! ML framework, enabling:
|
|
//!
|
|
//! - Multi-backend inference (CPU, WebGPU, CUDA, `PyTorch`)
|
|
//! - Model portability across different hardware
|
|
//! - ONNX model import
|
|
//! - Training with automatic differentiation
|
|
//!
|
|
//! ## Features
|
|
//!
|
|
//! - **ndarray**: CPU backend using ndarray (default)
|
|
//! - **wgpu**: WebGPU backend for cross-platform GPU support
|
|
//! - **cuda**: CUDA backend for NVIDIA GPUs
|
|
//! - **tch**: `PyTorch` backend via tch-rs
|
|
//!
|
|
//! ## Example
|
|
//!
|
|
//! ```ignore
|
|
//! use rtx_burn::{BurnSession, BurnConfig};
|
|
//! use rtx_tensor::{Tensor, Device};
|
|
//!
|
|
//! // Create session with WebGPU backend
|
|
//! let config = BurnConfig::default().with_backend(BurnBackend::Wgpu);
|
|
//! let session = BurnSession::new(config)?;
|
|
//!
|
|
//! // Load model and run inference
|
|
//! let model = session.load_model("model.onnx")?;
|
|
//! let input = Tensor::randn(vec![1, 3, 224, 224], &Device::CPU)?;
|
|
//! let output = session.run(&model, &input)?;
|
|
//! ```
|
|
|
|
#![allow(clippy::module_name_repetitions)]
|
|
#![allow(clippy::must_use_candidate)]
|
|
#![allow(clippy::missing_errors_doc)]
|
|
|
|
pub mod backend;
|
|
pub mod error;
|
|
pub mod model;
|
|
pub mod session;
|
|
pub mod tensor_bridge;
|
|
|
|
// Re-export main types
|
|
pub use backend::{BurnBackend, detect_best_backend};
|
|
pub use error::{BurnError, Result};
|
|
pub use model::{BurnModel, ModelInfo};
|
|
pub use session::{BurnConfig, BurnSession, SessionStats};
|
|
pub use tensor_bridge::{burn_to_rtx, rtx_to_burn};
|
|
|
|
/// Version information
|
|
pub const VERSION: &str = env!("CARGO_PKG_VERSION");
|
|
|
|
/// Check if running with WGPU backend available
|
|
pub fn is_wgpu_available() -> bool {
|
|
cfg!(feature = "wgpu")
|
|
}
|
|
|
|
/// Check if running with CUDA backend available
|
|
pub fn is_cuda_available() -> bool {
|
|
cfg!(feature = "cuda")
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
|
|
#[test]
|
|
fn test_version() {
|
|
assert!(!VERSION.is_empty());
|
|
}
|
|
|
|
#[test]
|
|
fn test_backend_detection() {
|
|
let backend = detect_best_backend();
|
|
// Should always return something
|
|
assert!(matches!(
|
|
backend,
|
|
BurnBackend::NdArray | BurnBackend::Wgpu | BurnBackend::Cuda | BurnBackend::Tch
|
|
));
|
|
}
|
|
}
|