Initial commit
This commit is contained in:
redclawsystems
@@ -0,0 +1,296 @@
|
||||
//! Burn inference session management
|
||||
//!
|
||||
//! Provides a high-level API for running inference with Burn models.
|
||||
|
||||
use crate::backend::{BackendConfig, BurnBackend};
|
||||
use crate::error::{BurnError, Result};
|
||||
use crate::model::BurnModel;
|
||||
use crate::tensor_bridge::{burn_to_rtx, rtx_to_burn};
|
||||
use rtx_tensor::{Device, Tensor};
|
||||
use serde::{Deserialize, Serialize};
|
||||
use std::collections::HashMap;
|
||||
use std::path::Path;
|
||||
use tracing::{debug, info};
|
||||
|
||||
/// Configuration for a Burn session
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct BurnConfig {
|
||||
/// Backend configuration
|
||||
pub backend: BackendConfig,
|
||||
/// Enable model caching
|
||||
pub cache_models: bool,
|
||||
/// Maximum cached models
|
||||
pub max_cached_models: usize,
|
||||
/// Output device for results
|
||||
pub output_device: Device,
|
||||
/// Enable profiling
|
||||
pub enable_profiling: bool,
|
||||
}
|
||||
|
||||
impl Default for BurnConfig {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
backend: BackendConfig::default(),
|
||||
cache_models: true,
|
||||
max_cached_models: 10,
|
||||
output_device: Device::Cpu,
|
||||
enable_profiling: false,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl BurnConfig {
|
||||
/// Set the backend
|
||||
pub fn with_backend(mut self, backend: BurnBackend) -> Self {
|
||||
self.backend.backend = backend;
|
||||
self
|
||||
}
|
||||
|
||||
/// Set the device index
|
||||
pub fn with_device(mut self, index: usize) -> Self {
|
||||
self.backend.device_index = index;
|
||||
self
|
||||
}
|
||||
|
||||
/// Set the output device
|
||||
pub fn with_output_device(mut self, device: Device) -> Self {
|
||||
self.output_device = device;
|
||||
self
|
||||
}
|
||||
|
||||
/// Enable profiling
|
||||
pub fn with_profiling(mut self) -> Self {
|
||||
self.enable_profiling = true;
|
||||
self
|
||||
}
|
||||
}
|
||||
|
||||
/// Statistics for a Burn session
|
||||
#[derive(Debug, Clone, Default)]
|
||||
pub struct SessionStats {
|
||||
/// Total inferences run
|
||||
pub total_inferences: u64,
|
||||
/// Total inference time in milliseconds
|
||||
pub total_inference_ms: f64,
|
||||
/// Models currently cached
|
||||
pub cached_models: usize,
|
||||
/// Cache hits
|
||||
pub cache_hits: u64,
|
||||
/// Cache misses
|
||||
pub cache_misses: u64,
|
||||
}
|
||||
|
||||
impl SessionStats {
|
||||
/// Average inference time
|
||||
pub fn avg_inference_ms(&self) -> f64 {
|
||||
if self.total_inferences == 0 {
|
||||
0.0
|
||||
} else {
|
||||
self.total_inference_ms / self.total_inferences as f64
|
||||
}
|
||||
}
|
||||
|
||||
/// Cache hit rate
|
||||
pub fn cache_hit_rate(&self) -> f64 {
|
||||
let total = self.cache_hits + self.cache_misses;
|
||||
if total == 0 {
|
||||
0.0
|
||||
} else {
|
||||
self.cache_hits as f64 / total as f64
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Burn inference session
|
||||
///
|
||||
/// Manages model loading, caching, and inference execution.
|
||||
pub struct BurnSession {
|
||||
/// Session configuration
|
||||
config: BurnConfig,
|
||||
/// Loaded models cache
|
||||
models: HashMap<String, BurnModel>,
|
||||
/// Session statistics
|
||||
stats: SessionStats,
|
||||
}
|
||||
|
||||
impl BurnSession {
|
||||
/// Create a new session with the given configuration
|
||||
pub fn new(config: BurnConfig) -> Result<Self> {
|
||||
info!(
|
||||
"Creating Burn session with {} backend",
|
||||
config.backend.backend.name()
|
||||
);
|
||||
|
||||
// Validate backend availability
|
||||
if !config.backend.backend.is_available() {
|
||||
return Err(BurnError::BackendUnavailable(
|
||||
config.backend.backend.name().to_string(),
|
||||
));
|
||||
}
|
||||
|
||||
// Initialize seed if provided
|
||||
if let Some(seed) = config.backend.seed {
|
||||
debug!("Setting seed: {}", seed);
|
||||
// Burn seed initialization would go here
|
||||
}
|
||||
|
||||
Ok(Self {
|
||||
config,
|
||||
models: HashMap::new(),
|
||||
stats: SessionStats::default(),
|
||||
})
|
||||
}
|
||||
|
||||
/// Create a session with default configuration
|
||||
pub fn default_session() -> Result<Self> {
|
||||
Self::new(BurnConfig::default())
|
||||
}
|
||||
|
||||
/// Load a model from a file path
|
||||
pub fn load_model(&mut self, path: impl AsRef<Path>) -> Result<&BurnModel> {
|
||||
let path = path.as_ref();
|
||||
let path_str = path.to_string_lossy().to_string();
|
||||
|
||||
// Check cache first
|
||||
if self.config.cache_models && self.models.contains_key(&path_str) {
|
||||
self.stats.cache_hits += 1;
|
||||
debug!("Model cache hit: {}", path_str);
|
||||
return Ok(self.models.get(&path_str).unwrap());
|
||||
}
|
||||
|
||||
self.stats.cache_misses += 1;
|
||||
|
||||
// Load the model
|
||||
info!("Loading model from: {}", path.display());
|
||||
let model = BurnModel::load(path, &self.config.backend)?;
|
||||
|
||||
// Cache if enabled
|
||||
if self.config.cache_models {
|
||||
// Evict oldest if at capacity
|
||||
if self.models.len() >= self.config.max_cached_models
|
||||
&& let Some(key) = self.models.keys().next().cloned() {
|
||||
debug!("Evicting cached model: {}", key);
|
||||
self.models.remove(&key);
|
||||
}
|
||||
self.models.insert(path_str.clone(), model);
|
||||
self.stats.cached_models = self.models.len();
|
||||
} else {
|
||||
self.models.insert(path_str.clone(), model);
|
||||
}
|
||||
|
||||
Ok(self.models.get(&path_str).unwrap())
|
||||
}
|
||||
|
||||
/// Run inference on a model
|
||||
pub fn run(
|
||||
&mut self,
|
||||
model: &BurnModel,
|
||||
inputs: HashMap<String, &Tensor>,
|
||||
) -> Result<HashMap<String, Tensor>> {
|
||||
let start = std::time::Instant::now();
|
||||
|
||||
// Convert inputs to Burn format
|
||||
let mut burn_inputs = HashMap::new();
|
||||
for (name, tensor) in inputs {
|
||||
let (data, shape) = rtx_to_burn(tensor)?;
|
||||
burn_inputs.insert(name.clone(), (data, shape));
|
||||
}
|
||||
|
||||
// Run inference (simulated for now - real implementation would use Burn's runtime)
|
||||
let burn_outputs = model.forward(burn_inputs)?;
|
||||
|
||||
// Convert outputs back to rtx tensors
|
||||
let mut outputs = HashMap::new();
|
||||
for (name, (data, shape)) in burn_outputs {
|
||||
let tensor = burn_to_rtx(data, shape, &self.config.output_device)?;
|
||||
outputs.insert(name, tensor);
|
||||
}
|
||||
|
||||
// Update stats
|
||||
let elapsed = start.elapsed().as_secs_f64() * 1000.0;
|
||||
self.stats.total_inferences += 1;
|
||||
self.stats.total_inference_ms += elapsed;
|
||||
|
||||
debug!("Inference completed in {:.2}ms", elapsed);
|
||||
Ok(outputs)
|
||||
}
|
||||
|
||||
/// Run inference with a single input/output
|
||||
pub fn run_simple(&mut self, model: &BurnModel, input: &Tensor) -> Result<Tensor> {
|
||||
let inputs = HashMap::from([("input".to_string(), input)]);
|
||||
let mut outputs = self.run(model, inputs)?;
|
||||
|
||||
outputs
|
||||
.remove("output")
|
||||
.ok_or_else(|| BurnError::Inference("No output tensor found".to_string()))
|
||||
}
|
||||
|
||||
/// Get session configuration
|
||||
pub fn config(&self) -> &BurnConfig {
|
||||
&self.config
|
||||
}
|
||||
|
||||
/// Get session statistics
|
||||
pub fn stats(&self) -> &SessionStats {
|
||||
&self.stats
|
||||
}
|
||||
|
||||
/// Clear the model cache
|
||||
pub fn clear_cache(&mut self) {
|
||||
self.models.clear();
|
||||
self.stats.cached_models = 0;
|
||||
info!("Model cache cleared");
|
||||
}
|
||||
|
||||
/// Get the backend being used
|
||||
pub fn backend(&self) -> BurnBackend {
|
||||
self.config.backend.backend
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn test_config_default() {
|
||||
let config = BurnConfig::default();
|
||||
assert!(config.cache_models);
|
||||
assert_eq!(config.max_cached_models, 10);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_config_builder() {
|
||||
let config = BurnConfig::default()
|
||||
.with_backend(BurnBackend::NdArray)
|
||||
.with_device(1)
|
||||
.with_profiling();
|
||||
|
||||
assert_eq!(config.backend.backend, BurnBackend::NdArray);
|
||||
assert_eq!(config.backend.device_index, 1);
|
||||
assert!(config.enable_profiling);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_session_stats() {
|
||||
let mut stats = SessionStats::default();
|
||||
stats.total_inferences = 10;
|
||||
stats.total_inference_ms = 100.0;
|
||||
stats.cache_hits = 8;
|
||||
stats.cache_misses = 2;
|
||||
|
||||
assert!((stats.avg_inference_ms() - 10.0).abs() < 0.001);
|
||||
assert!((stats.cache_hit_rate() - 0.8).abs() < 0.001);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_session_creation() {
|
||||
// NdArray backend should always be available with the default feature
|
||||
let config = BurnConfig::default().with_backend(BurnBackend::NdArray);
|
||||
let session = BurnSession::new(config);
|
||||
// May fail if ndarray feature is not enabled, that's ok for test
|
||||
if let Ok(session) = session {
|
||||
assert_eq!(session.backend(), BurnBackend::NdArray);
|
||||
}
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user