Initial commit
This commit is contained in:
redclawsystems
@@ -0,0 +1,472 @@
|
||||
//! Early stopping strategies with rtx-validation integration
|
||||
//!
|
||||
//! This module provides early stopping functionality for AutoML optimization,
|
||||
//! including cross-validation based stopping using rtx-validation's KFold.
|
||||
|
||||
use crate::{AutoMLError, AutoMLResult};
|
||||
use rtx_tensor::Tensor;
|
||||
use rtx_validation::cv::{CrossValidator, KFold};
|
||||
use serde::{Deserialize, Serialize};
|
||||
use std::collections::HashMap;
|
||||
use std::time::Instant;
|
||||
|
||||
/// Early stopping controller for AutoML optimization
|
||||
pub struct EarlyStopping {
|
||||
criteria: Option<StoppingCriteria>,
|
||||
trials: Vec<TrialResult>,
|
||||
cv_trials: Vec<CVTrialResult>,
|
||||
start_time: Option<Instant>,
|
||||
best_score: f64,
|
||||
best_trial_id: Option<usize>,
|
||||
trials_since_improvement: usize,
|
||||
cv_folds: usize,
|
||||
multi_objectives: Vec<String>,
|
||||
pareto_front: Vec<TrialResult>,
|
||||
adaptive_patience_enabled: bool,
|
||||
}
|
||||
|
||||
impl EarlyStopping {
|
||||
/// Create a new EarlyStopping instance
|
||||
pub fn new() -> AutoMLResult<Self> {
|
||||
Ok(Self {
|
||||
criteria: None,
|
||||
trials: Vec::new(),
|
||||
cv_trials: Vec::new(),
|
||||
start_time: None,
|
||||
best_score: f64::NEG_INFINITY,
|
||||
best_trial_id: None,
|
||||
trials_since_improvement: 0,
|
||||
cv_folds: 5,
|
||||
multi_objectives: Vec::new(),
|
||||
pareto_front: Vec::new(),
|
||||
adaptive_patience_enabled: false,
|
||||
})
|
||||
}
|
||||
|
||||
/// Set the stopping criteria
|
||||
pub fn set_criteria(&mut self, criteria: StoppingCriteria) {
|
||||
self.criteria = Some(criteria);
|
||||
}
|
||||
|
||||
/// Record a trial result
|
||||
pub fn record_trial(&mut self, trial: TrialResult) {
|
||||
let score = trial.score;
|
||||
let trial_id = trial.trial_id;
|
||||
self.trials.push(trial);
|
||||
|
||||
// Check for improvement
|
||||
if let Some(ref criteria) = self.criteria {
|
||||
let improvement = score - self.best_score;
|
||||
if improvement > criteria.min_improvement {
|
||||
self.best_score = score;
|
||||
self.best_trial_id = Some(trial_id);
|
||||
self.trials_since_improvement = 0;
|
||||
} else {
|
||||
self.trials_since_improvement += 1;
|
||||
}
|
||||
} else if score > self.best_score {
|
||||
self.best_score = score;
|
||||
self.best_trial_id = Some(trial_id);
|
||||
self.trials_since_improvement = 0;
|
||||
} else {
|
||||
self.trials_since_improvement += 1;
|
||||
}
|
||||
}
|
||||
|
||||
/// Check if optimization should stop
|
||||
pub fn should_stop(&self) -> bool {
|
||||
if let Some(ref criteria) = self.criteria {
|
||||
// Check patience
|
||||
if self.trials_since_improvement >= criteria.patience {
|
||||
return true;
|
||||
}
|
||||
|
||||
// Check max trials
|
||||
if self.trials.len() >= criteria.max_trials {
|
||||
return true;
|
||||
}
|
||||
|
||||
// Check target score
|
||||
if let Some(target) = criteria.target_score
|
||||
&& self.best_score >= target
|
||||
{
|
||||
return true;
|
||||
}
|
||||
|
||||
// Check time budget
|
||||
if let Some(start) = self.start_time
|
||||
&& start.elapsed().as_secs() >= criteria.max_time_seconds
|
||||
{
|
||||
return true;
|
||||
}
|
||||
}
|
||||
|
||||
false
|
||||
}
|
||||
|
||||
/// Start the optimization timer
|
||||
pub fn start_timer(&mut self) {
|
||||
self.start_time = Some(Instant::now());
|
||||
}
|
||||
|
||||
/// Get remaining time in seconds
|
||||
pub fn get_remaining_time_seconds(&self) -> f64 {
|
||||
if let Some(ref criteria) = self.criteria {
|
||||
if let Some(start) = self.start_time {
|
||||
let elapsed = start.elapsed().as_secs_f64();
|
||||
let budget = criteria.max_time_seconds as f64;
|
||||
return (budget - elapsed).max(0.0);
|
||||
}
|
||||
return criteria.max_time_seconds as f64;
|
||||
}
|
||||
3600.0 // Default 1 hour
|
||||
}
|
||||
|
||||
/// Get the optimization history
|
||||
pub fn get_optimization_history(&self) -> OptimizationHistory {
|
||||
let total_time = self
|
||||
.start_time
|
||||
.map_or(0.0, |s| s.elapsed().as_secs_f64());
|
||||
|
||||
OptimizationHistory {
|
||||
trials: self.trials.clone(),
|
||||
best_score: self.best_score,
|
||||
best_trial_id: self.best_trial_id,
|
||||
total_time_seconds: total_time,
|
||||
}
|
||||
}
|
||||
|
||||
/// Get current patience (trials without improvement)
|
||||
pub fn get_current_patience(&self) -> usize {
|
||||
self.trials_since_improvement
|
||||
}
|
||||
|
||||
/// Enable adaptive patience based on learning dynamics
|
||||
pub fn enable_adaptive_patience(&mut self, enabled: bool) {
|
||||
self.adaptive_patience_enabled = enabled;
|
||||
}
|
||||
|
||||
/// Enable multi-objective optimization
|
||||
pub fn enable_multi_objective(&mut self, objectives: Vec<String>) {
|
||||
self.multi_objectives = objectives;
|
||||
}
|
||||
|
||||
/// Record a multi-objective trial
|
||||
pub fn record_multi_objective_trial(
|
||||
&mut self,
|
||||
trial_id: usize,
|
||||
objectives: HashMap<String, f64>,
|
||||
) {
|
||||
// Create trial result from objectives
|
||||
let score = objectives.values().sum::<f64>() / objectives.len() as f64;
|
||||
let trial = TrialResult {
|
||||
trial_id,
|
||||
score,
|
||||
training_time_seconds: 0.0,
|
||||
memory_usage_mb: 0.0,
|
||||
hyperparameters: HashMap::new(),
|
||||
model_name: String::new(),
|
||||
};
|
||||
|
||||
// Update Pareto front
|
||||
self.update_pareto_front(trial.clone(), &objectives);
|
||||
self.trials.push(trial);
|
||||
}
|
||||
|
||||
fn update_pareto_front(&mut self, trial: TrialResult, _objectives: &HashMap<String, f64>) {
|
||||
// Simplified Pareto dominance check
|
||||
// In a full implementation, this would check dominance across all objectives
|
||||
let dominated = self.pareto_front.iter().any(|p| p.score > trial.score);
|
||||
if !dominated {
|
||||
self.pareto_front.retain(|p| p.score >= trial.score);
|
||||
self.pareto_front.push(trial);
|
||||
}
|
||||
}
|
||||
|
||||
/// Check if multi-objective optimization should stop
|
||||
pub fn should_stop_multi_objective(&self) -> AutoMLResult<bool> {
|
||||
// Stop if Pareto front hasn't changed recently
|
||||
if let Some(ref criteria) = self.criteria
|
||||
&& self.trials_since_improvement >= criteria.patience
|
||||
{
|
||||
return Ok(true);
|
||||
}
|
||||
Ok(false)
|
||||
}
|
||||
|
||||
/// Get the Pareto front of solutions
|
||||
pub fn get_pareto_front(&self) -> Vec<TrialResult> {
|
||||
self.pareto_front.clone()
|
||||
}
|
||||
|
||||
// ============================================================
|
||||
// CV-BASED STOPPING WITH rtx-validation INTEGRATION
|
||||
// ============================================================
|
||||
|
||||
/// Enable CV-based stopping with KFold from rtx-validation
|
||||
pub fn enable_cv_based_stopping(&mut self, folds: usize) {
|
||||
self.cv_folds = folds;
|
||||
}
|
||||
|
||||
/// Record a CV trial result with proper validation
|
||||
pub fn record_cv_trial(&mut self, trial_id: usize, cv_scores: Vec<f64>, mean: f64, std: f64) {
|
||||
let cv_result = CVTrialResult {
|
||||
trial_id,
|
||||
cv_scores,
|
||||
mean_score: mean,
|
||||
std_score: std,
|
||||
};
|
||||
self.cv_trials.push(cv_result);
|
||||
|
||||
// Update best score with CV confidence adjustment
|
||||
// Use lower bound estimate (mean - std) for conservative comparison
|
||||
let confidence_adjusted_score = mean - std;
|
||||
if confidence_adjusted_score > self.best_score {
|
||||
self.best_score = confidence_adjusted_score;
|
||||
self.best_trial_id = Some(trial_id);
|
||||
self.trials_since_improvement = 0;
|
||||
} else {
|
||||
self.trials_since_improvement += 1;
|
||||
}
|
||||
}
|
||||
|
||||
/// Check if should stop based on CV confidence intervals
|
||||
pub fn should_stop_with_cv_confidence(&self) -> AutoMLResult<bool> {
|
||||
if let Some(ref criteria) = self.criteria {
|
||||
// Check patience with CV-adjusted scores
|
||||
if self.trials_since_improvement >= criteria.patience {
|
||||
return Ok(true);
|
||||
}
|
||||
|
||||
// Check if recent trials show no statistically significant improvement
|
||||
if self.cv_trials.len() >= 3 {
|
||||
let recent_trials: Vec<_> = self.cv_trials.iter().rev().take(3).collect();
|
||||
let all_overlap = recent_trials.windows(2).all(|pair| {
|
||||
let (t1, t2) = (&pair[0], &pair[1]);
|
||||
// Check if confidence intervals overlap (no significant improvement)
|
||||
let t1_upper = t1.mean_score + t1.std_score;
|
||||
let t2_lower = t2.mean_score - t2.std_score;
|
||||
t1_upper >= t2_lower
|
||||
});
|
||||
|
||||
if all_overlap && self.cv_trials.len() >= criteria.min_trials {
|
||||
return Ok(true);
|
||||
}
|
||||
}
|
||||
}
|
||||
Ok(false)
|
||||
}
|
||||
|
||||
/// Get CV history for analysis
|
||||
pub fn get_cv_history(&self) -> Vec<CVTrialResult> {
|
||||
self.cv_trials.clone()
|
||||
}
|
||||
|
||||
/// Perform CV evaluation using rtx-validation KFold
|
||||
///
|
||||
/// This method creates KFold splits and evaluates using the provided function.
|
||||
/// Returns (mean_score, std_score, fold_scores).
|
||||
pub fn evaluate_with_cv<F>(
|
||||
&self,
|
||||
x: &Tensor,
|
||||
y: &Tensor,
|
||||
mut eval_fn: F,
|
||||
) -> AutoMLResult<(f64, f64, Vec<f64>)>
|
||||
where
|
||||
F: FnMut(&[usize], &[usize]) -> AutoMLResult<f64>,
|
||||
{
|
||||
let cv = KFold::new(self.cv_folds).shuffle(true).random_state(42);
|
||||
let splits = cv
|
||||
.split(x, Some(y))
|
||||
.map_err(|e| AutoMLError::ValidationError(e.to_string()))?;
|
||||
|
||||
let mut fold_scores = Vec::new();
|
||||
for (train_idx, val_idx) in splits {
|
||||
let score = eval_fn(&train_idx, &val_idx)?;
|
||||
fold_scores.push(score);
|
||||
}
|
||||
|
||||
let mean = fold_scores.iter().sum::<f64>() / fold_scores.len() as f64;
|
||||
let variance =
|
||||
fold_scores.iter().map(|s| (s - mean).powi(2)).sum::<f64>() / fold_scores.len() as f64;
|
||||
let std = variance.sqrt();
|
||||
|
||||
Ok((mean, std, fold_scores))
|
||||
}
|
||||
|
||||
/// Serialize to JSON
|
||||
pub fn to_json(&self) -> AutoMLResult<String> {
|
||||
let history = self.get_optimization_history();
|
||||
serde_json::to_string(&history).map_err(|e| AutoMLError::SerializationError(e.to_string()))
|
||||
}
|
||||
|
||||
/// Deserialize from JSON (creates new instance with history)
|
||||
pub fn from_json(_json: &str) -> AutoMLResult<Self> {
|
||||
// For now, return a fresh instance
|
||||
// Full implementation would restore state from JSON
|
||||
Self::new()
|
||||
}
|
||||
}
|
||||
|
||||
impl Default for EarlyStopping {
|
||||
fn default() -> Self {
|
||||
Self::new().unwrap()
|
||||
}
|
||||
}
|
||||
|
||||
/// Stopping criteria configuration
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct StoppingCriteria {
|
||||
/// Number of trials without improvement before stopping
|
||||
pub patience: usize,
|
||||
/// Minimum improvement to reset patience
|
||||
pub min_improvement: f64,
|
||||
/// Maximum number of trials
|
||||
pub max_trials: usize,
|
||||
/// Maximum time in seconds
|
||||
pub max_time_seconds: u64,
|
||||
/// Target score to achieve (optional)
|
||||
pub target_score: Option<f64>,
|
||||
/// Minimum trials before stopping
|
||||
pub min_trials: usize,
|
||||
}
|
||||
|
||||
impl Default for StoppingCriteria {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
patience: 10,
|
||||
min_improvement: 0.001,
|
||||
max_trials: 100,
|
||||
max_time_seconds: 3600,
|
||||
target_score: None,
|
||||
min_trials: 5,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Result of a single trial
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct TrialResult {
|
||||
pub trial_id: usize,
|
||||
pub score: f64,
|
||||
pub training_time_seconds: f64,
|
||||
pub memory_usage_mb: f64,
|
||||
pub hyperparameters: HashMap<String, String>,
|
||||
pub model_name: String,
|
||||
}
|
||||
|
||||
/// Optimization history
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct OptimizationHistory {
|
||||
pub trials: Vec<TrialResult>,
|
||||
pub best_score: f64,
|
||||
pub best_trial_id: Option<usize>,
|
||||
pub total_time_seconds: f64,
|
||||
}
|
||||
|
||||
/// CV trial result with statistical information
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct CVTrialResult {
|
||||
pub trial_id: usize,
|
||||
pub cv_scores: Vec<f64>,
|
||||
pub mean_score: f64,
|
||||
pub std_score: f64,
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn test_early_stopping_creation() {
|
||||
let es = EarlyStopping::new().unwrap();
|
||||
assert_eq!(es.get_current_patience(), 0);
|
||||
assert!(!es.should_stop());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_trial_recording() {
|
||||
let mut es = EarlyStopping::new().unwrap();
|
||||
es.set_criteria(StoppingCriteria {
|
||||
patience: 3,
|
||||
min_improvement: 0.01,
|
||||
..Default::default()
|
||||
});
|
||||
|
||||
// Record improving trials
|
||||
es.record_trial(TrialResult {
|
||||
trial_id: 0,
|
||||
score: 0.5,
|
||||
training_time_seconds: 1.0,
|
||||
memory_usage_mb: 100.0,
|
||||
hyperparameters: HashMap::new(),
|
||||
model_name: "test".to_string(),
|
||||
});
|
||||
assert_eq!(es.best_score, 0.5);
|
||||
assert_eq!(es.get_current_patience(), 0);
|
||||
|
||||
// Record non-improving trial
|
||||
es.record_trial(TrialResult {
|
||||
trial_id: 1,
|
||||
score: 0.5,
|
||||
training_time_seconds: 1.0,
|
||||
memory_usage_mb: 100.0,
|
||||
hyperparameters: HashMap::new(),
|
||||
model_name: "test".to_string(),
|
||||
});
|
||||
assert_eq!(es.get_current_patience(), 1);
|
||||
assert!(!es.should_stop());
|
||||
|
||||
// Record more non-improving trials to trigger stopping
|
||||
es.record_trial(TrialResult {
|
||||
trial_id: 2,
|
||||
score: 0.5,
|
||||
..Default::default()
|
||||
});
|
||||
es.record_trial(TrialResult {
|
||||
trial_id: 3,
|
||||
score: 0.5,
|
||||
..Default::default()
|
||||
});
|
||||
assert!(es.should_stop());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_cv_based_stopping() {
|
||||
let mut es = EarlyStopping::new().unwrap();
|
||||
es.enable_cv_based_stopping(5);
|
||||
es.set_criteria(StoppingCriteria {
|
||||
patience: 2,
|
||||
min_trials: 3,
|
||||
..Default::default()
|
||||
});
|
||||
|
||||
// Record CV trials
|
||||
es.record_cv_trial(0, vec![0.8, 0.82, 0.79, 0.81, 0.80], 0.804, 0.011);
|
||||
assert!(!es.should_stop_with_cv_confidence().unwrap());
|
||||
|
||||
es.record_cv_trial(1, vec![0.81, 0.83, 0.80, 0.82, 0.81], 0.814, 0.012);
|
||||
assert!(!es.should_stop_with_cv_confidence().unwrap());
|
||||
|
||||
// Stagnating performance
|
||||
es.record_cv_trial(2, vec![0.81, 0.82, 0.80, 0.81, 0.80], 0.808, 0.008);
|
||||
es.record_cv_trial(3, vec![0.80, 0.81, 0.79, 0.80, 0.80], 0.800, 0.007);
|
||||
|
||||
// Should detect overlapping confidence intervals
|
||||
let history = es.get_cv_history();
|
||||
assert_eq!(history.len(), 4);
|
||||
}
|
||||
}
|
||||
|
||||
impl Default for TrialResult {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
trial_id: 0,
|
||||
score: 0.0,
|
||||
training_time_seconds: 0.0,
|
||||
memory_usage_mb: 0.0,
|
||||
hyperparameters: HashMap::new(),
|
||||
model_name: String::new(),
|
||||
}
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user