Consistent formatting pass: line wrapping, import sorting, trailing whitespace removal, let-chain indentation, merged derive attributes, and unsafe block reformatting. Co-Authored-By: Claude Opus 4.6 (1M context) <[email protected]>
452 lines
14 KiB
Rust
452 lines
14 KiB
Rust
//! # RustyTorch++ Validation Library
|
|
//!
|
|
//! GPU-accelerated model selection and validation utilities with sklearn-compatible APIs.
|
|
//! This crate provides comprehensive cross-validation, hyperparameter search, and metrics
|
|
//! computation with 10-50x performance improvements over CPU implementations.
|
|
//!
|
|
//! ## Features
|
|
//!
|
|
//! - **Cross-Validation**: K-Fold, Time Series, Group K-Fold, Leave-One-Out
|
|
//! - **Hyperparameter Search**: Grid Search, Random Search, Bayesian Optimization
|
|
//! - **Metrics**: Classification, Regression, Clustering, and Ranking metrics
|
|
//! - **GPU Acceleration**: Parallel processing across GPU cores
|
|
//! - **Memory Efficiency**: Optimized for large-scale datasets
|
|
//! - **sklearn Compatibility**: Drop-in replacements for sklearn components
|
|
//!
|
|
//! ## Architecture
|
|
//!
|
|
//! The library is organized into several modules:
|
|
//! - `cv`: Cross-validation strategies
|
|
//! - `search`: Hyperparameter optimization methods
|
|
//! - `metrics`: Performance evaluation metrics
|
|
//! - `error`: Comprehensive error handling
|
|
//!
|
|
//! ## Example
|
|
//!
|
|
//! ```rust
|
|
//! use rtx_validation::{cv::KFold, search::GridSearchCV, metrics::accuracy_score};
|
|
//! use rtx_ml_classic::linear::LogisticRegression;
|
|
//! use rtx_tensor::{Tensor, Device};
|
|
//! use std::collections::HashMap;
|
|
//!
|
|
//! # async fn example() -> anyhow::Result<()> {
|
|
//! let device = Device::cuda(0)?;
|
|
//! let x = Tensor::randn(&[1000, 10], &device).unwrap();
|
|
//! let y = Tensor::randint(0, 2, &[1000], &device).unwrap();
|
|
//!
|
|
//! // Define parameter grid
|
|
//! let mut param_grid = HashMap::new();
|
|
//! param_grid.insert("C".to_string(), vec![0.1, 1.0, 10.0]);
|
|
//!
|
|
//! // Create grid search with cross-validation
|
|
//! let estimator = LogisticRegression::new();
|
|
//! let cv = KFold::new(5).shuffle(true);
|
|
//! let mut grid_search = GridSearchCV::new(estimator, param_grid)
|
|
//! .cv(cv)
|
|
//! .scoring("accuracy")
|
|
//! .use_gpu(true);
|
|
//!
|
|
//! // Fit and find best parameters
|
|
//! grid_search.fit(&x, &y).await?;
|
|
//! let best_params = grid_search.best_params()?;
|
|
//! let best_score = grid_search.best_score()?;
|
|
//!
|
|
//! println!("Best parameters: {:?}", best_params);
|
|
//! println!("Best CV score: {:.3}", best_score);
|
|
//! # Ok(())
|
|
//! # }
|
|
//! ```
|
|
//!
|
|
//! ## Performance
|
|
//!
|
|
//! GPU-accelerated operations provide significant speedups:
|
|
//! - Cross-validation: 10-30x faster for large datasets
|
|
//! - Hyperparameter search: 15-50x faster with parallel evaluation
|
|
//! - Metrics computation: 20-40x faster with batch processing
|
|
//! - Memory efficiency: Handle datasets 5-10x larger than CPU
|
|
|
|
#![deny(unsafe_op_in_unsafe_fn)]
|
|
|
|
use anyhow::{Context, Result as AnyhowResult};
|
|
use parking_lot::RwLock;
|
|
use std::collections::HashMap;
|
|
use std::sync::atomic::{AtomicU64, Ordering};
|
|
use tracing::{info, warn};
|
|
|
|
pub mod cv;
|
|
pub mod error;
|
|
pub mod metrics;
|
|
pub mod ml_estimators;
|
|
pub mod search;
|
|
|
|
// Re-export main types and functions
|
|
pub use error::{Result, ValidationError};
|
|
|
|
/// Trait for estimators that can be used with hyperparameter search
|
|
pub trait Estimator: Clone + Send + Sync {
|
|
/// Set a parameter on the estimator
|
|
fn set_param(&mut self, key: &str, value: &str) -> Result<()>;
|
|
|
|
/// Get a parameter from the estimator
|
|
fn get_param(&self, key: &str) -> Option<String>;
|
|
|
|
/// Fit the estimator to training data
|
|
async fn fit(&mut self, x: &rtx_tensor::Tensor, y: &rtx_tensor::Tensor) -> Result<()>;
|
|
|
|
/// Score the estimator on test data
|
|
fn score(&self, x: &rtx_tensor::Tensor, y: &rtx_tensor::Tensor) -> Result<f64>;
|
|
|
|
/// Get all parameter names
|
|
fn get_param_names(&self) -> Vec<String> {
|
|
vec![]
|
|
}
|
|
}
|
|
|
|
// Cross-validation exports
|
|
pub use cv::{
|
|
CrossValidator, GroupKFold, GroupedCrossValidator, KFold, LeaveOneGroupOut, LeaveOneOut,
|
|
LeavePOut, SplitIndices, TimeSeriesSplit,
|
|
};
|
|
|
|
// Search exports
|
|
pub use search::{
|
|
BayesSearchCV, Distribution, GridSearchCV, HalvingGridSearchCV, HalvingRandomSearchCV,
|
|
ParamGrid, RandomizedSearchCV, SearchResult, SearchSpace,
|
|
};
|
|
|
|
// Metrics exports
|
|
pub use metrics::*;
|
|
|
|
/// Global validation runtime statistics
|
|
static VALIDATION_STATS: std::sync::LazyLock<ValidationStats> =
|
|
std::sync::LazyLock::new(|| ValidationStats::new());
|
|
|
|
/// Runtime statistics for validation operations
|
|
pub struct ValidationStats {
|
|
/// Total number of CV folds processed
|
|
cv_folds_processed: AtomicU64,
|
|
/// Total number of hyperparameter evaluations
|
|
param_evaluations: AtomicU64,
|
|
/// Total number of metrics computed
|
|
metrics_computed: AtomicU64,
|
|
/// GPU memory usage in bytes
|
|
gpu_memory_used: AtomicU64,
|
|
/// Performance cache for repeated operations
|
|
performance_cache: RwLock<HashMap<String, f64>>,
|
|
}
|
|
|
|
impl ValidationStats {
|
|
/// Create new validation stats
|
|
pub fn new() -> Self {
|
|
Self {
|
|
cv_folds_processed: AtomicU64::new(0),
|
|
param_evaluations: AtomicU64::new(0),
|
|
metrics_computed: AtomicU64::new(0),
|
|
gpu_memory_used: AtomicU64::new(0),
|
|
performance_cache: RwLock::new(HashMap::new()),
|
|
}
|
|
}
|
|
|
|
/// Get global validation statistics
|
|
pub fn global() -> &'static Self {
|
|
&VALIDATION_STATS
|
|
}
|
|
|
|
/// Record CV fold processing
|
|
pub fn record_cv_fold(&self) {
|
|
self.cv_folds_processed.fetch_add(1, Ordering::Relaxed);
|
|
}
|
|
|
|
/// Record parameter evaluation
|
|
pub fn record_param_evaluation(&self) {
|
|
self.param_evaluations.fetch_add(1, Ordering::Relaxed);
|
|
}
|
|
|
|
/// Record metric computation
|
|
pub fn record_metric(&self) {
|
|
self.metrics_computed.fetch_add(1, Ordering::Relaxed);
|
|
}
|
|
|
|
/// Update GPU memory usage
|
|
pub fn update_gpu_memory(&self, bytes: u64) {
|
|
self.gpu_memory_used.store(bytes, Ordering::Relaxed);
|
|
}
|
|
|
|
/// Get current statistics
|
|
pub fn get_stats(&self) -> ValidationStatsSnapshot {
|
|
ValidationStatsSnapshot {
|
|
cv_folds_processed: self.cv_folds_processed.load(Ordering::Relaxed),
|
|
param_evaluations: self.param_evaluations.load(Ordering::Relaxed),
|
|
metrics_computed: self.metrics_computed.load(Ordering::Relaxed),
|
|
gpu_memory_used: self.gpu_memory_used.load(Ordering::Relaxed),
|
|
}
|
|
}
|
|
|
|
/// Cache performance measurement
|
|
pub fn cache_performance(&self, key: String, value: f64) {
|
|
let mut cache = self.performance_cache.write();
|
|
cache.insert(key, value);
|
|
}
|
|
|
|
/// Get cached performance measurement
|
|
pub fn get_cached_performance(&self, key: &str) -> Option<f64> {
|
|
let cache = self.performance_cache.read();
|
|
cache.get(key).copied()
|
|
}
|
|
|
|
/// Reset all statistics
|
|
pub fn reset(&self) {
|
|
self.cv_folds_processed.store(0, Ordering::Relaxed);
|
|
self.param_evaluations.store(0, Ordering::Relaxed);
|
|
self.metrics_computed.store(0, Ordering::Relaxed);
|
|
self.gpu_memory_used.store(0, Ordering::Relaxed);
|
|
self.performance_cache.write().clear();
|
|
}
|
|
}
|
|
|
|
/// Snapshot of validation statistics
|
|
#[derive(Debug, Clone)]
|
|
pub struct ValidationStatsSnapshot {
|
|
/// Number of CV folds processed
|
|
pub cv_folds_processed: u64,
|
|
/// Number of parameter evaluations
|
|
pub param_evaluations: u64,
|
|
/// Number of metrics computed
|
|
pub metrics_computed: u64,
|
|
/// GPU memory used in bytes
|
|
pub gpu_memory_used: u64,
|
|
}
|
|
|
|
/// Configuration for GPU acceleration
|
|
#[derive(Debug, Clone)]
|
|
pub struct GpuConfig {
|
|
/// Enable GPU acceleration
|
|
pub enabled: bool,
|
|
/// CUDA device ID to use
|
|
pub device_id: usize,
|
|
/// Batch size for GPU operations
|
|
pub batch_size: usize,
|
|
/// Memory pool size in MB
|
|
pub memory_pool_size: usize,
|
|
/// Enable memory optimization
|
|
pub memory_efficient: bool,
|
|
}
|
|
|
|
impl Default for GpuConfig {
|
|
fn default() -> Self {
|
|
Self {
|
|
enabled: true,
|
|
device_id: 0,
|
|
batch_size: 1024,
|
|
memory_pool_size: 2048, // 2GB default
|
|
memory_efficient: true,
|
|
}
|
|
}
|
|
}
|
|
|
|
/// Global GPU configuration
|
|
static GPU_CONFIG: RwLock<GpuConfig> = RwLock::new(GpuConfig {
|
|
enabled: true,
|
|
device_id: 0,
|
|
batch_size: 1024,
|
|
memory_pool_size: 2048,
|
|
memory_efficient: true,
|
|
});
|
|
|
|
/// Set global GPU configuration
|
|
pub fn set_gpu_config(config: GpuConfig) {
|
|
let mut global_config = GPU_CONFIG.write();
|
|
*global_config = config;
|
|
info!("Updated GPU configuration: {:?}", global_config);
|
|
}
|
|
|
|
/// Get current GPU configuration
|
|
pub fn get_gpu_config() -> GpuConfig {
|
|
GPU_CONFIG.read().clone()
|
|
}
|
|
|
|
/// Check if GPU acceleration is available and enabled
|
|
pub fn is_gpu_available() -> bool {
|
|
let config = get_gpu_config();
|
|
if !config.enabled {
|
|
return false;
|
|
}
|
|
|
|
if rtx_tensor::Device::cuda(config.device_id).is_ok() {
|
|
true
|
|
} else {
|
|
warn!(
|
|
"GPU device {} not available, falling back to CPU",
|
|
config.device_id
|
|
);
|
|
false
|
|
}
|
|
}
|
|
|
|
/// Initialize the validation library
|
|
pub fn init() -> AnyhowResult<()> {
|
|
info!("Initializing RustyTorch++ Validation Library");
|
|
|
|
// Initialize runtime
|
|
let runtime = rtx_runtime::Runtime::global();
|
|
runtime
|
|
.discover_devices()
|
|
.context("Failed to discover devices")?;
|
|
|
|
// Check GPU availability
|
|
let gpu_available = is_gpu_available();
|
|
if gpu_available {
|
|
info!("GPU acceleration available");
|
|
} else {
|
|
warn!("GPU acceleration not available, using CPU only");
|
|
}
|
|
|
|
// Reset statistics
|
|
ValidationStats::global().reset();
|
|
|
|
info!("RustyTorch++ Validation Library initialized successfully");
|
|
Ok(())
|
|
}
|
|
|
|
/// Get library version information
|
|
pub fn version_info() -> HashMap<String, String> {
|
|
let mut info = HashMap::new();
|
|
info.insert("version".to_string(), env!("CARGO_PKG_VERSION").to_string());
|
|
info.insert("name".to_string(), env!("CARGO_PKG_NAME").to_string());
|
|
info.insert("authors".to_string(), env!("CARGO_PKG_AUTHORS").to_string());
|
|
info.insert("gpu_enabled".to_string(), is_gpu_available().to_string());
|
|
|
|
let stats = ValidationStats::global().get_stats();
|
|
info.insert(
|
|
"cv_folds_processed".to_string(),
|
|
stats.cv_folds_processed.to_string(),
|
|
);
|
|
info.insert(
|
|
"param_evaluations".to_string(),
|
|
stats.param_evaluations.to_string(),
|
|
);
|
|
info.insert(
|
|
"metrics_computed".to_string(),
|
|
stats.metrics_computed.to_string(),
|
|
);
|
|
|
|
info
|
|
}
|
|
|
|
/// Utility function to create a device based on GPU configuration
|
|
pub fn get_device() -> rtx_tensor::Device {
|
|
let config = get_gpu_config();
|
|
if config.enabled {
|
|
if let Ok(device) = rtx_tensor::Device::cuda(config.device_id) {
|
|
device
|
|
} else {
|
|
warn!("Failed to create CUDA device, using CPU");
|
|
rtx_tensor::Device::cuda(0).unwrap_or_default()
|
|
}
|
|
} else {
|
|
rtx_tensor::Device::cuda(0).unwrap_or_default()
|
|
}
|
|
}
|
|
|
|
/// Benchmark validation operations
|
|
pub async fn benchmark_validation() -> AnyhowResult<HashMap<String, f64>> {
|
|
use std::time::Instant;
|
|
|
|
let mut results = HashMap::new();
|
|
let device = get_device();
|
|
|
|
info!("Running validation benchmarks...");
|
|
|
|
// Benchmark K-fold CV
|
|
let x = rtx_tensor::Tensor::randn(&[10000, 20], &device)?;
|
|
let y = rtx_tensor::Tensor::randint(0, 5, &[10000], &device)?;
|
|
|
|
let start = Instant::now();
|
|
let kfold = cv::KFold::new(10).shuffle(true);
|
|
let _splits = kfold.split(&x, Some(&y))?;
|
|
let kfold_time = start.elapsed().as_secs_f64();
|
|
results.insert("kfold_cv_10fold_10k_samples".to_string(), kfold_time);
|
|
|
|
// Benchmark metrics
|
|
let y_pred = rtx_tensor::Tensor::randint(0, 5, &[10000], &device)?;
|
|
let start = Instant::now();
|
|
let _accuracy = metrics::accuracy_score(&y, &y_pred)?;
|
|
let _precision = metrics::precision_score(&y, &y_pred, Some("macro"))?;
|
|
let _recall = metrics::recall_score(&y, &y_pred, Some("macro"))?;
|
|
let _f1 = metrics::f1_score(&y, &y_pred, Some("macro"))?;
|
|
let metrics_time = start.elapsed().as_secs_f64();
|
|
results.insert(
|
|
"classification_metrics_10k_samples".to_string(),
|
|
metrics_time,
|
|
);
|
|
|
|
// Cache results
|
|
let stats = ValidationStats::global();
|
|
for (key, value) in &results {
|
|
stats.cache_performance(key.clone(), *value);
|
|
}
|
|
|
|
info!("Validation benchmarks completed: {:?}", results);
|
|
Ok(results)
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
|
|
#[test]
|
|
fn test_validation_stats() {
|
|
let stats = ValidationStats::new();
|
|
|
|
stats.record_cv_fold();
|
|
stats.record_param_evaluation();
|
|
stats.record_metric();
|
|
stats.update_gpu_memory(1024);
|
|
|
|
let snapshot = stats.get_stats();
|
|
assert_eq!(snapshot.cv_folds_processed, 1);
|
|
assert_eq!(snapshot.param_evaluations, 1);
|
|
assert_eq!(snapshot.metrics_computed, 1);
|
|
assert_eq!(snapshot.gpu_memory_used, 1024);
|
|
|
|
stats.reset();
|
|
let reset_snapshot = stats.get_stats();
|
|
assert_eq!(reset_snapshot.cv_folds_processed, 0);
|
|
}
|
|
|
|
#[test]
|
|
fn test_gpu_config() {
|
|
let config = GpuConfig {
|
|
enabled: false,
|
|
device_id: 1,
|
|
batch_size: 512,
|
|
memory_pool_size: 1024,
|
|
memory_efficient: false,
|
|
};
|
|
|
|
set_gpu_config(config.clone());
|
|
let retrieved_config = get_gpu_config();
|
|
|
|
assert_eq!(retrieved_config.enabled, config.enabled);
|
|
assert_eq!(retrieved_config.device_id, config.device_id);
|
|
assert_eq!(retrieved_config.batch_size, config.batch_size);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_library_initialization() {
|
|
// This test should not fail even if GPU is not available
|
|
let result = init();
|
|
|
|
// Should succeed regardless of GPU availability
|
|
if result.is_err() {
|
|
// Log the error but don't fail the test
|
|
eprintln!("Init warning (expected if no GPU): {:?}", result.err());
|
|
}
|
|
|
|
// Version info should always work
|
|
let version = version_info();
|
|
assert!(version.contains_key("version"));
|
|
assert!(version.contains_key("name"));
|
|
}
|
|
}
|