7ef92983f9
Mechanical auto-fixes: redundant borrows, clone on Copy, or_insert_with, single-char push_str, get(0) → first(), needless borrow, let_and_return. 150 files, no behavior changes. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
1424 lines
45 KiB
Rust
1424 lines
45 KiB
Rust
//! Early Stopping Infrastructure for Hyperparameter Optimization
|
||
//!
|
||
//! This module provides production-ready early stopping strategies for hyperparameter
|
||
//! optimization trials. It implements multiple stopping criteria including:
|
||
//!
|
||
//! - **Plateau Detection**: Stop when loss plateaus (no improvement for N epochs)
|
||
//! - **Median Pruner**: Stop trials performing worse than median
|
||
//! - **Percentile Pruner**: Stop trials in bottom N percentile
|
||
//! - **Observer Pattern**: Extensible callback system for custom strategies
|
||
//!
|
||
//! ## Design Principles
|
||
//!
|
||
//! 1. **Orthogonal**: Works independently of optimizer internals
|
||
//! 2. **Zero-overhead**: Disabled by default, no cost when not used
|
||
//! 3. **Type-safe**: Generic over parameter types
|
||
//! 4. **Production-ready**: Comprehensive error handling and logging
|
||
//!
|
||
//! ## Architecture
|
||
//!
|
||
//! ```text
|
||
//! ┌─────────────────────────────────────────────────────────────┐
|
||
//! │ TrialObserver Trait │
|
||
//! │ ┌──────────────────┐ ┌────────────────────────────┐ │
|
||
//! │ │ on_trial_start │ │ on_epoch_complete │ │
|
||
//! │ └──────────────────┘ └────────────────────────────┘ │
|
||
//! └─────────────────────────────────────────────────────────────┘
|
||
//! ↓
|
||
//! ┌─────────────────────────────┐
|
||
//! │ EarlyStoppingObserver │
|
||
//! │ ┌─────────────────────┐ │
|
||
//! │ │ EarlyStoppingConfig │ │
|
||
//! │ └─────────────────────┘ │
|
||
//! │ ┌─────────────────────┐ │
|
||
//! │ │ Per-trial State │ │
|
||
//! │ │ HashMap<id, State> │ │
|
||
//! │ └─────────────────────┘ │
|
||
//! └─────────────────────────────┘
|
||
//! ↓
|
||
//! ┌─────────────────────────────┐
|
||
//! │ EarlyStoppingStrategy │
|
||
//! │ ┌───────────────────────┐ │
|
||
//! │ │ Plateau │ │
|
||
//! │ │ MedianPruner │ │
|
||
//! │ │ PercentilePruner │ │
|
||
//! │ │ SuccessiveHalving │ │
|
||
//! │ │ Hyperband │ │
|
||
//! │ └───────────────────────┘ │
|
||
//! └─────────────────────────────┘
|
||
//! ```
|
||
//!
|
||
//! ## Usage Example
|
||
//!
|
||
//! ```rust,no_run
|
||
//! use ml::hyperopt::early_stopping::{
|
||
//! EarlyStoppingConfig, EarlyStoppingObserver, EarlyStoppingStrategy,
|
||
//! TrialObserver, ObserverDecision, EpochMetrics,
|
||
//! };
|
||
//!
|
||
//! // Create observer with plateau detection
|
||
//! let config = EarlyStoppingConfig {
|
||
//! patience_epochs: 10,
|
||
//! min_delta: 1e-4,
|
||
//! strategy: EarlyStoppingStrategy::Plateau,
|
||
//! ..Default::default()
|
||
//! };
|
||
//! let mut observer = EarlyStoppingObserver::new(config);
|
||
//!
|
||
//! // In training loop
|
||
//! observer.on_trial_start(0, "learning_rate=0.001");
|
||
//!
|
||
//! for epoch in 0..100 {
|
||
//! // ... training code ...
|
||
//! let val_loss = 0.5; // From validation
|
||
//!
|
||
//! let metrics = EpochMetrics {
|
||
//! epoch,
|
||
//! train_loss: 0.4,
|
||
//! val_loss,
|
||
//! timestamp: epoch as f64,
|
||
//! };
|
||
//!
|
||
//! match observer.on_epoch_complete(0, epoch, &metrics) {
|
||
//! ObserverDecision::Continue => continue,
|
||
//! ObserverDecision::StopTrial => {
|
||
//! println!("Trial stopped early at epoch {}", epoch);
|
||
//! break;
|
||
//! }
|
||
//! ObserverDecision::StopStudy => {
|
||
//! println!("Study stopped early");
|
||
//! return;
|
||
//! }
|
||
//! }
|
||
//! }
|
||
//!
|
||
//! observer.on_trial_complete(0, 0.5);
|
||
//! ```
|
||
//!
|
||
//! ## Performance Characteristics
|
||
//!
|
||
//! - **Memory**: O(trials × epochs) for history tracking
|
||
//! - **Per-epoch overhead**: ~1-10μs for decision logic
|
||
//! - **Cross-trial pruning**: O(trials) for median computation
|
||
//!
|
||
//! ## Integration with Adapters
|
||
//!
|
||
//! Adapters can integrate early stopping by:
|
||
//!
|
||
//! 1. Creating observer: `let observer = EarlyStoppingObserver::new(config)`
|
||
//! 2. Calling `on_trial_start()` before training
|
||
//! 3. Calling `on_epoch_complete()` after each epoch
|
||
//! 4. Breaking training loop on `StopTrial` decision
|
||
//! 5. Calling `on_trial_complete()` or `on_trial_failed()` at end
|
||
|
||
use serde::{Deserialize, Serialize};
|
||
use std::collections::HashMap;
|
||
use tracing::{info, warn};
|
||
|
||
/// Configuration for early stopping criteria
|
||
///
|
||
/// Defines when to stop training trials based on validation loss behavior.
|
||
/// All configurations support serialization for experiment tracking.
|
||
///
|
||
/// ## Parameter Guidelines
|
||
///
|
||
/// - **patience_epochs**: 5-20 epochs depending on model convergence speed
|
||
/// - **min_delta**: 1e-4 for stable models, 1e-3 for noisy training
|
||
/// - **min_epochs**: 20-50 epochs to avoid premature stopping
|
||
/// - **validation_frequency**: 1 for small models, 2-5 for large models
|
||
///
|
||
/// ## Example
|
||
///
|
||
/// ```rust
|
||
/// use ml::hyperopt::early_stopping::{EarlyStoppingConfig, EarlyStoppingStrategy};
|
||
///
|
||
/// let config = EarlyStoppingConfig {
|
||
/// patience_epochs: 15,
|
||
/// min_delta: 5e-4,
|
||
/// min_epochs: 30,
|
||
/// strategy: EarlyStoppingStrategy::Plateau,
|
||
/// ..Default::default()
|
||
/// };
|
||
/// ```
|
||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||
pub struct EarlyStoppingConfig {
|
||
/// Number of epochs with no improvement before stopping
|
||
///
|
||
/// Training stops if validation loss doesn't improve by at least
|
||
/// `min_delta` for this many consecutive epochs.
|
||
pub patience_epochs: usize,
|
||
|
||
/// Minimum improvement to reset patience (absolute value)
|
||
///
|
||
/// An improvement is considered significant if:
|
||
/// `old_loss - new_loss > min_delta`
|
||
pub min_delta: f64,
|
||
|
||
/// Whether to compare to best trial across all trials
|
||
///
|
||
/// When true, uses study-wide best loss as baseline instead of
|
||
/// trial's own best loss. Useful for aggressive pruning.
|
||
pub compare_to_baseline: bool,
|
||
|
||
/// Frequency of validation checks (every N epochs)
|
||
///
|
||
/// Set to 1 for epoch-by-epoch monitoring, or higher values
|
||
/// to reduce overhead for expensive validation.
|
||
pub validation_frequency: usize,
|
||
|
||
/// Minimum epochs before early stopping can trigger
|
||
///
|
||
/// Prevents premature stopping during initial exploration.
|
||
/// Recommended: 20-50 epochs depending on model complexity.
|
||
pub min_epochs: usize,
|
||
|
||
/// Strategy to use for early stopping decision
|
||
///
|
||
/// Available strategies:
|
||
/// - `Plateau`: Loss plateau detection (single-trial)
|
||
/// - `MedianPruner`: Prune if worse than median (cross-trial)
|
||
/// - `PercentilePruner`: Prune bottom N% (cross-trial)
|
||
/// - `SuccessiveHalving`: Aggressive halving (cross-trial)
|
||
/// - `Hyperband`: Optimal resource allocation (cross-trial)
|
||
pub strategy: EarlyStoppingStrategy,
|
||
}
|
||
|
||
impl Default for EarlyStoppingConfig {
|
||
fn default() -> Self {
|
||
Self {
|
||
patience_epochs: 10,
|
||
min_delta: 1e-4,
|
||
compare_to_baseline: false,
|
||
validation_frequency: 1,
|
||
min_epochs: 20,
|
||
strategy: EarlyStoppingStrategy::Plateau,
|
||
}
|
||
}
|
||
}
|
||
|
||
/// State of a trial in the optimization process
|
||
///
|
||
/// Tracks the lifecycle of a single trial from initialization through
|
||
/// completion or early termination. Used for result analysis and debugging.
|
||
///
|
||
/// ## State Transitions
|
||
///
|
||
/// ```text
|
||
/// Pending → Running → {EarlyStopped*, Completed, Failed}
|
||
/// ↓
|
||
/// Pruned*
|
||
/// ```
|
||
///
|
||
/// *Early termination states preserve the epoch where stopping occurred
|
||
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
|
||
pub enum TrialState {
|
||
/// Trial is waiting to start
|
||
Pending,
|
||
|
||
/// Trial is currently running
|
||
Running {
|
||
/// Current epoch number
|
||
current_epoch: usize,
|
||
},
|
||
|
||
/// Trial stopped early due to loss plateau
|
||
///
|
||
/// Triggered when validation loss fails to improve for
|
||
/// `patience_epochs` consecutive epochs.
|
||
EarlyStoppedPlateau {
|
||
/// Epoch where stopping occurred
|
||
stopped_at_epoch: usize,
|
||
},
|
||
|
||
/// Trial stopped early due to threshold violation
|
||
///
|
||
/// Triggered when validation loss exceeds a predefined
|
||
/// maximum threshold (e.g., divergence detection).
|
||
EarlyStoppedThreshold {
|
||
/// Epoch where stopping occurred
|
||
stopped_at_epoch: usize,
|
||
},
|
||
|
||
/// Trial stopped early by pruner (median/percentile)
|
||
///
|
||
/// Triggered by cross-trial comparison strategies when trial
|
||
/// performs worse than the study median or target percentile.
|
||
Pruned {
|
||
/// Epoch where pruning occurred
|
||
stopped_at_epoch: usize,
|
||
/// Human-readable reason (e.g., "Worse than median")
|
||
reason: String,
|
||
},
|
||
|
||
/// Trial completed successfully
|
||
Completed,
|
||
|
||
/// Trial failed with error
|
||
Failed {
|
||
/// Error message describing failure cause
|
||
error: String,
|
||
},
|
||
}
|
||
|
||
/// Per-trial early stopping state
|
||
///
|
||
/// Maintains all state needed to make early stopping decisions for a single trial.
|
||
/// Includes best loss tracking, patience counter, and epoch history for analysis.
|
||
///
|
||
/// ## Memory Usage
|
||
///
|
||
/// - Fixed: ~48 bytes (best_loss, counters, flags)
|
||
/// - Variable: ~48 bytes × epochs (epoch history)
|
||
/// - Total: ~48 + 48N bytes for N epochs
|
||
///
|
||
/// ## Example
|
||
///
|
||
/// ```rust
|
||
/// use ml::hyperopt::early_stopping::EarlyStoppingState;
|
||
///
|
||
/// let mut state = EarlyStoppingState::new();
|
||
///
|
||
/// // Update with validation loss
|
||
/// let improved = state.update(0.5, 1e-4);
|
||
/// assert!(improved); // First update always improves
|
||
///
|
||
/// let improved = state.update(0.3, 1e-4);
|
||
/// assert!(improved); // Significant improvement
|
||
///
|
||
/// let improved = state.update(0.299, 1e-4);
|
||
/// assert!(!improved); // Marginal improvement (< min_delta)
|
||
/// ```
|
||
#[derive(Debug, Clone)]
|
||
pub struct EarlyStoppingState {
|
||
/// Best validation loss seen so far
|
||
pub best_val_loss: f64,
|
||
|
||
/// Consecutive epochs without improvement
|
||
pub patience_counter: usize,
|
||
|
||
/// Whether trial was stopped early
|
||
pub stopped: bool,
|
||
|
||
/// Epoch where stopping occurred (if stopped)
|
||
pub stopped_at_epoch: Option<usize>,
|
||
|
||
/// Complete epoch history for this trial
|
||
pub epoch_history: Vec<EpochMetrics>,
|
||
}
|
||
|
||
impl EarlyStoppingState {
|
||
/// Create a new early stopping state
|
||
///
|
||
/// Initializes with infinite loss (no observations yet) and
|
||
/// zero patience counter.
|
||
pub fn new() -> Self {
|
||
Self {
|
||
best_val_loss: f64::INFINITY,
|
||
patience_counter: 0,
|
||
stopped: false,
|
||
stopped_at_epoch: None,
|
||
epoch_history: Vec::new(),
|
||
}
|
||
}
|
||
|
||
/// Update state with new validation loss
|
||
///
|
||
/// Compares new loss to best loss and updates patience counter.
|
||
/// Returns true if this represents a significant improvement.
|
||
///
|
||
/// # Arguments
|
||
///
|
||
/// * `val_loss` - Current validation loss
|
||
/// * `min_delta` - Minimum improvement threshold
|
||
///
|
||
/// # Returns
|
||
///
|
||
/// True if `best_val_loss - val_loss > min_delta`, false otherwise
|
||
///
|
||
/// # Example
|
||
///
|
||
/// ```rust
|
||
/// # use ml::hyperopt::early_stopping::EarlyStoppingState;
|
||
/// let mut state = EarlyStoppingState::new();
|
||
///
|
||
/// assert!(state.update(0.5, 1e-4)); // First update
|
||
/// assert!(state.update(0.3, 1e-4)); // Big improvement
|
||
/// assert!(!state.update(0.299, 1e-4)); // Marginal (< 1e-4)
|
||
/// ```
|
||
pub fn update(&mut self, val_loss: f64, min_delta: f64) -> bool {
|
||
// Check if this is a significant improvement
|
||
let improvement = self.best_val_loss - val_loss;
|
||
|
||
// Use epsilon comparison to avoid floating point precision issues
|
||
// e.g., 0.5 - 0.49 = 0.010000000000000009 (not exactly 0.01)
|
||
const EPSILON: f64 = 1e-10;
|
||
if improvement > min_delta + EPSILON {
|
||
// Significant improvement - update best and reset patience
|
||
self.best_val_loss = val_loss;
|
||
self.patience_counter = 0;
|
||
true
|
||
} else {
|
||
// No significant improvement - increment patience
|
||
self.patience_counter += 1;
|
||
false
|
||
}
|
||
}
|
||
|
||
/// Record epoch metrics for history tracking
|
||
///
|
||
/// Stores metrics for later analysis, convergence plots, and debugging.
|
||
///
|
||
/// # Arguments
|
||
///
|
||
/// * `metrics` - Epoch metrics to record
|
||
pub fn record_epoch_metrics(&mut self, metrics: EpochMetrics) {
|
||
self.epoch_history.push(metrics);
|
||
}
|
||
|
||
/// Mark trial as stopped early
|
||
///
|
||
/// # Arguments
|
||
///
|
||
/// * `epoch` - Epoch where stopping occurred
|
||
pub fn mark_stopped(&mut self, epoch: usize) {
|
||
self.stopped = true;
|
||
self.stopped_at_epoch = Some(epoch);
|
||
}
|
||
}
|
||
|
||
impl Default for EarlyStoppingState {
|
||
fn default() -> Self {
|
||
Self::new()
|
||
}
|
||
}
|
||
|
||
/// Metrics for a single training epoch
|
||
///
|
||
/// Contains all metrics needed for early stopping decisions and
|
||
/// result analysis. Serializable for experiment tracking.
|
||
///
|
||
/// ## Usage
|
||
///
|
||
/// Create after each epoch and pass to `TrialObserver::on_epoch_complete()`:
|
||
///
|
||
/// ```rust
|
||
/// use ml::hyperopt::early_stopping::EpochMetrics;
|
||
/// use std::time::{SystemTime, UNIX_EPOCH};
|
||
///
|
||
/// let metrics = EpochMetrics {
|
||
/// epoch: 10,
|
||
/// train_loss: 0.45,
|
||
/// val_loss: 0.52,
|
||
/// timestamp: SystemTime::now()
|
||
/// .duration_since(UNIX_EPOCH)
|
||
/// .unwrap()
|
||
/// .as_secs_f64(),
|
||
/// };
|
||
/// ```
|
||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||
pub struct EpochMetrics {
|
||
/// Epoch number (0-indexed)
|
||
pub epoch: usize,
|
||
|
||
/// Training loss for this epoch
|
||
pub train_loss: f64,
|
||
|
||
/// Validation loss for this epoch (used for early stopping decisions)
|
||
pub val_loss: f64,
|
||
|
||
/// Unix timestamp when epoch completed
|
||
pub timestamp: f64,
|
||
}
|
||
|
||
/// Available early stopping strategies
|
||
///
|
||
/// Each strategy implements a different approach to trial termination:
|
||
///
|
||
/// - **Single-trial**: Decisions based only on current trial (Plateau)
|
||
/// - **Cross-trial**: Decisions based on comparison with other trials (Median, Percentile, etc.)
|
||
///
|
||
/// ## Strategy Comparison
|
||
///
|
||
/// | Strategy | Type | Warmup | Aggressiveness | Use Case |
|
||
/// |----------|------|--------|----------------|----------|
|
||
/// | Plateau | Single | None | Low | Stable convergence |
|
||
/// | MedianPruner | Cross | Required | Medium | Balanced exploration |
|
||
/// | PercentilePruner | Cross | Required | High | Aggressive pruning |
|
||
/// | SuccessiveHalving | Cross | None | Very High | Known good ranges |
|
||
/// | Hyperband | Cross | None | Optimal | Large-scale search |
|
||
///
|
||
/// ## Example
|
||
///
|
||
/// ```rust
|
||
/// use ml::hyperopt::early_stopping::EarlyStoppingStrategy;
|
||
///
|
||
/// // Conservative - only stop on clear plateau
|
||
/// let plateau = EarlyStoppingStrategy::Plateau;
|
||
///
|
||
/// // Balanced - prune bottom half after warmup
|
||
/// let median = EarlyStoppingStrategy::MedianPruner { warmup_steps: 5 };
|
||
///
|
||
/// // Aggressive - prune bottom 25% after warmup
|
||
/// let percentile = EarlyStoppingStrategy::PercentilePruner {
|
||
/// percentile: 25.0,
|
||
/// warmup_steps: 10,
|
||
/// };
|
||
/// ```
|
||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||
pub enum EarlyStoppingStrategy {
|
||
/// Stop if no improvement for N epochs (single-trial)
|
||
///
|
||
/// Most conservative strategy. Only considers current trial's
|
||
/// own performance history. Good for:
|
||
/// - Small-scale optimization (< 30 trials)
|
||
/// - Unstable training dynamics
|
||
/// - When you want every trial to reach natural convergence
|
||
Plateau,
|
||
|
||
/// Stop if worse than median of all trials (cross-trial)
|
||
///
|
||
/// Compares current trial's loss to the median loss across all
|
||
/// completed and running trials at the same epoch. Prunes trials
|
||
/// in the bottom 50%.
|
||
///
|
||
/// **Parameters**:
|
||
/// - `warmup_steps`: Minimum epochs before pruning can occur
|
||
///
|
||
/// **Best for**: Balanced exploration vs. exploitation
|
||
MedianPruner {
|
||
/// Minimum epochs before pruning can trigger
|
||
warmup_steps: usize,
|
||
},
|
||
|
||
/// Stop if in bottom N percentile (cross-trial)
|
||
///
|
||
/// More aggressive than median pruning. Prunes trials performing
|
||
/// worse than the Nth percentile across all trials.
|
||
///
|
||
/// **Parameters**:
|
||
/// - `percentile`: Target percentile (0-100). Lower = more aggressive
|
||
/// - `warmup_steps`: Minimum epochs before pruning
|
||
///
|
||
/// **Example**: `percentile=25.0` prunes bottom 25% of trials
|
||
///
|
||
/// **Best for**: Large-scale search (100+ trials) with tight budgets
|
||
PercentilePruner {
|
||
/// Target percentile threshold (0-100)
|
||
percentile: f64,
|
||
/// Minimum epochs before pruning
|
||
warmup_steps: usize,
|
||
},
|
||
|
||
/// Successive halving algorithm (cross-trial)
|
||
///
|
||
/// Aggressively halves the number of trials at each bracket.
|
||
/// Allocates resources to best-performing trials only.
|
||
///
|
||
/// **Parameters**:
|
||
/// - `reduction_factor`: How many trials to keep (default: 2 = keep half)
|
||
///
|
||
/// **Best for**: Known good parameter ranges, want to zoom in fast
|
||
SuccessiveHalving {
|
||
/// Reduction factor (2 = keep half, 3 = keep third, etc.)
|
||
reduction_factor: usize,
|
||
},
|
||
|
||
/// Hyperband algorithm (cross-trial)
|
||
///
|
||
/// Runs multiple successive halving brackets with different
|
||
/// resource allocations. Provably optimal for any smooth function.
|
||
///
|
||
/// **Parameters**:
|
||
/// - `max_resource`: Maximum epochs per trial
|
||
/// - `reduction_factor`: Halving rate (typically 2-4)
|
||
///
|
||
/// **Best for**: Large-scale search with unknown convergence behavior
|
||
Hyperband {
|
||
/// Maximum resource allocation (epochs)
|
||
max_resource: usize,
|
||
/// Reduction factor for successive halving
|
||
reduction_factor: usize,
|
||
},
|
||
}
|
||
|
||
/// Trait for early stopping strategies
|
||
///
|
||
/// Implement this trait to create custom stopping logic. The trait is
|
||
/// called by `EarlyStoppingObserver` after each epoch to make decisions.
|
||
///
|
||
/// ## Contract
|
||
///
|
||
/// - Must be deterministic (same inputs → same output)
|
||
/// - Should be fast (called every epoch for every trial)
|
||
/// - Can access cross-trial history for comparison-based pruning
|
||
///
|
||
/// ## Example Implementation
|
||
///
|
||
/// ```rust,ignore
|
||
/// struct CustomStrategy {
|
||
/// threshold: f64,
|
||
/// }
|
||
///
|
||
/// impl EarlyStoppingStrategyTrait for CustomStrategy {
|
||
/// fn should_stop(
|
||
/// &self,
|
||
/// epoch: usize,
|
||
/// val_loss: f64,
|
||
/// state: &mut EarlyStoppingState,
|
||
/// trial_history: &[f64],
|
||
/// ) -> bool {
|
||
/// val_loss > self.threshold
|
||
/// }
|
||
///
|
||
/// fn name(&self) -> &'static str {
|
||
/// "custom_threshold"
|
||
/// }
|
||
/// }
|
||
/// ```
|
||
pub trait EarlyStoppingStrategyTrait {
|
||
/// Check if trial should stop
|
||
///
|
||
/// # Arguments
|
||
///
|
||
/// * `epoch` - Current epoch number
|
||
/// * `val_loss` - Current validation loss
|
||
/// * `state` - Mutable trial state for updates
|
||
/// * `trial_history` - Best losses from other trials (for cross-trial strategies)
|
||
///
|
||
/// # Returns
|
||
///
|
||
/// True if trial should stop, false to continue
|
||
fn should_stop(
|
||
&self,
|
||
epoch: usize,
|
||
val_loss: f64,
|
||
state: &mut EarlyStoppingState,
|
||
trial_history: &[f64],
|
||
) -> bool;
|
||
|
||
/// Name of the strategy for logging
|
||
fn name(&self) -> &'static str;
|
||
}
|
||
|
||
/// Plateau detection strategy implementation
|
||
///
|
||
/// Stops training when validation loss plateaus for `patience_epochs`
|
||
/// consecutive epochs without improvement >= `min_delta`.
|
||
///
|
||
/// ## Algorithm
|
||
///
|
||
/// ```text
|
||
/// for each epoch:
|
||
/// if val_loss improves by >= min_delta:
|
||
/// reset patience_counter = 0
|
||
/// else:
|
||
/// patience_counter += 1
|
||
///
|
||
/// if patience_counter >= patience_epochs:
|
||
/// return STOP
|
||
/// ```
|
||
///
|
||
/// ## Performance
|
||
///
|
||
/// - Time: O(1) per epoch
|
||
/// - Memory: O(1)
|
||
/// - Overhead: ~1μs per decision
|
||
#[derive(Debug, Clone)]
|
||
pub struct PlateauDetectionStrategy {
|
||
patience_epochs: usize,
|
||
min_delta: f64,
|
||
}
|
||
|
||
impl PlateauDetectionStrategy {
|
||
/// Create new plateau detection strategy
|
||
///
|
||
/// # Arguments
|
||
///
|
||
/// * `patience_epochs` - Number of epochs to wait for improvement
|
||
/// * `min_delta` - Minimum improvement threshold
|
||
pub fn new(patience_epochs: usize, min_delta: f64) -> Self {
|
||
Self {
|
||
patience_epochs,
|
||
min_delta,
|
||
}
|
||
}
|
||
}
|
||
|
||
impl EarlyStoppingStrategyTrait for PlateauDetectionStrategy {
|
||
fn should_stop(
|
||
&self,
|
||
_epoch: usize,
|
||
val_loss: f64,
|
||
state: &mut EarlyStoppingState,
|
||
_trial_history: &[f64],
|
||
) -> bool {
|
||
// Update state and check if improved
|
||
state.update(val_loss, self.min_delta);
|
||
|
||
// Stop if patience exhausted
|
||
state.patience_counter >= self.patience_epochs
|
||
}
|
||
|
||
fn name(&self) -> &'static str {
|
||
"plateau_detection"
|
||
}
|
||
}
|
||
|
||
/// Median pruner strategy implementation
|
||
///
|
||
/// Prunes trials performing worse than the median across all trials
|
||
/// at the same epoch. Requires warmup period to accumulate statistics.
|
||
///
|
||
/// ## Algorithm
|
||
///
|
||
/// ```text
|
||
/// if epoch < warmup_steps:
|
||
/// return CONTINUE
|
||
///
|
||
/// median = compute_median(trial_history)
|
||
/// if val_loss > median:
|
||
/// return STOP
|
||
/// else:
|
||
/// return CONTINUE
|
||
/// ```
|
||
///
|
||
/// ## Performance
|
||
///
|
||
/// - Time: O(T log T) per decision, where T = number of trials
|
||
/// - Memory: O(1)
|
||
/// - Overhead: ~10-50μs per decision (depends on trial count)
|
||
///
|
||
/// ## Example
|
||
///
|
||
/// ```rust
|
||
/// use ml::hyperopt::early_stopping::MedianPrunerStrategy;
|
||
///
|
||
/// let pruner = MedianPrunerStrategy::new(5);
|
||
/// let trial_history = vec![0.3, 0.4, 0.5, 0.6, 0.7];
|
||
///
|
||
/// // Current trial worse than median (0.5)
|
||
/// assert!(pruner.should_stop(10, 0.8, &trial_history));
|
||
///
|
||
/// // Current trial better than median
|
||
/// assert!(!pruner.should_stop(10, 0.2, &trial_history));
|
||
/// ```
|
||
#[derive(Debug, Clone)]
|
||
pub struct MedianPrunerStrategy {
|
||
warmup_steps: usize,
|
||
}
|
||
|
||
impl MedianPrunerStrategy {
|
||
/// Create new median pruner strategy
|
||
///
|
||
/// # Arguments
|
||
///
|
||
/// * `warmup_steps` - Minimum epochs before pruning can occur
|
||
pub fn new(warmup_steps: usize) -> Self {
|
||
Self { warmup_steps }
|
||
}
|
||
|
||
/// Check if trial should stop based on median comparison
|
||
///
|
||
/// # Arguments
|
||
///
|
||
/// * `epoch` - Current epoch number
|
||
/// * `val_loss` - Current validation loss
|
||
/// * `trial_history` - Best losses from other trials
|
||
///
|
||
/// # Returns
|
||
///
|
||
/// True if current loss is worse than median, false otherwise
|
||
pub fn should_stop(&self, epoch: usize, val_loss: f64, trial_history: &[f64]) -> bool {
|
||
// Wait for warmup
|
||
if epoch < self.warmup_steps {
|
||
return false;
|
||
}
|
||
|
||
// Need at least 2 trials for meaningful comparison
|
||
if trial_history.len() < 2 {
|
||
return false;
|
||
}
|
||
|
||
// Compute median
|
||
let mut sorted = trial_history.to_vec();
|
||
sorted.sort_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal));
|
||
let median = if sorted.len() % 2 == 0 {
|
||
let mid = sorted.len() / 2;
|
||
(sorted[mid - 1] + sorted[mid]) / 2.0
|
||
} else {
|
||
sorted[sorted.len() / 2]
|
||
};
|
||
|
||
// Prune if worse than median
|
||
val_loss > median
|
||
}
|
||
}
|
||
|
||
/// Percentile pruner strategy implementation
|
||
///
|
||
/// Prunes trials in the bottom N percentile across all trials.
|
||
/// More aggressive than median pruning.
|
||
///
|
||
/// ## Algorithm
|
||
///
|
||
/// ```text
|
||
/// if epoch < warmup_steps:
|
||
/// return CONTINUE
|
||
///
|
||
/// threshold = compute_percentile(trial_history, percentile)
|
||
/// if val_loss > threshold:
|
||
/// return STOP
|
||
/// else:
|
||
/// return CONTINUE
|
||
/// ```
|
||
///
|
||
/// ## Performance
|
||
///
|
||
/// - Time: O(T log T) per decision
|
||
/// - Memory: O(1)
|
||
/// - Overhead: ~10-50μs per decision
|
||
///
|
||
/// ## Example
|
||
///
|
||
/// ```rust
|
||
/// use ml::hyperopt::early_stopping::PercentilePrunerStrategy;
|
||
///
|
||
/// // Prune bottom 25%
|
||
/// let pruner = PercentilePrunerStrategy::new(25.0, 5);
|
||
/// let trial_history = vec![0.2, 0.4, 0.6, 0.8];
|
||
///
|
||
/// // 25th percentile ≈ 0.35
|
||
/// assert!(pruner.should_stop(10, 0.9, &trial_history)); // Bottom 25%
|
||
/// assert!(!pruner.should_stop(10, 0.5, &trial_history)); // Above 25th
|
||
/// ```
|
||
#[derive(Debug, Clone)]
|
||
pub struct PercentilePrunerStrategy {
|
||
percentile: f64,
|
||
warmup_steps: usize,
|
||
}
|
||
|
||
impl PercentilePrunerStrategy {
|
||
/// Create new percentile pruner strategy
|
||
///
|
||
/// # Arguments
|
||
///
|
||
/// * `percentile` - Target percentile (0-100). Lower = more aggressive
|
||
/// * `warmup_steps` - Minimum epochs before pruning
|
||
///
|
||
/// # Panics
|
||
///
|
||
/// Panics if `percentile` is not in range [0, 100]
|
||
pub fn new(percentile: f64, warmup_steps: usize) -> Self {
|
||
assert!(
|
||
(0.0..=100.0).contains(&percentile),
|
||
"Percentile must be in [0, 100]"
|
||
);
|
||
Self {
|
||
percentile,
|
||
warmup_steps,
|
||
}
|
||
}
|
||
|
||
/// Check if trial should stop based on percentile comparison
|
||
///
|
||
/// # Arguments
|
||
///
|
||
/// * `epoch` - Current epoch number
|
||
/// * `val_loss` - Current validation loss
|
||
/// * `trial_history` - Best losses from other trials
|
||
///
|
||
/// # Returns
|
||
///
|
||
/// True if current loss is in bottom percentile, false otherwise
|
||
pub fn should_stop(&self, epoch: usize, val_loss: f64, trial_history: &[f64]) -> bool {
|
||
// Wait for warmup
|
||
if epoch < self.warmup_steps {
|
||
return false;
|
||
}
|
||
|
||
// Need sufficient trials for meaningful percentile
|
||
if trial_history.len() < 3 {
|
||
return false;
|
||
}
|
||
|
||
// Compute percentile threshold
|
||
let mut sorted = trial_history.to_vec();
|
||
sorted.sort_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal));
|
||
|
||
let index = ((self.percentile / 100.0) * (sorted.len() as f64 - 1.0)).round() as usize;
|
||
let threshold = sorted[index.min(sorted.len() - 1)];
|
||
|
||
// Prune if worse than threshold
|
||
val_loss > threshold
|
||
}
|
||
}
|
||
|
||
/// Decision from observer callbacks
|
||
///
|
||
/// Returned by `TrialObserver::on_epoch_complete()` to control trial execution.
|
||
///
|
||
/// ## Decision Flow
|
||
///
|
||
/// ```text
|
||
/// on_epoch_complete() → ObserverDecision
|
||
/// ↓
|
||
/// ┌──────────────────────┼──────────────────────┐
|
||
/// ↓ ↓ ↓
|
||
/// Continue StopTrial StopStudy
|
||
/// (keep going) (stop this trial) (stop all trials)
|
||
/// ```
|
||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||
pub enum ObserverDecision {
|
||
/// Continue training this trial
|
||
Continue,
|
||
|
||
/// Stop this trial early (save resources)
|
||
StopTrial,
|
||
|
||
/// Stop entire optimization study (found optimal or budget exhausted)
|
||
StopStudy,
|
||
}
|
||
|
||
/// Trait for trial lifecycle observers
|
||
///
|
||
/// Implement this trait to receive callbacks during trial execution.
|
||
/// Used by adapters to integrate early stopping logic.
|
||
///
|
||
/// ## Callback Order
|
||
///
|
||
/// ```text
|
||
/// on_trial_start(0)
|
||
/// ↓
|
||
/// on_epoch_complete(0, 0) → Continue
|
||
/// ↓
|
||
/// on_epoch_complete(0, 1) → Continue
|
||
/// ↓
|
||
/// on_epoch_complete(0, 2) → StopTrial
|
||
/// ↓
|
||
/// on_trial_complete(0, final_loss)
|
||
/// ```
|
||
///
|
||
/// ## Example Implementation
|
||
///
|
||
/// See `EarlyStoppingObserver` for full implementation.
|
||
pub trait TrialObserver {
|
||
/// Called when trial starts
|
||
///
|
||
/// # Arguments
|
||
///
|
||
/// * `trial_num` - Trial number (0-indexed)
|
||
/// * `params` - Human-readable parameter string
|
||
fn on_trial_start(&mut self, trial_num: usize, params: &str);
|
||
|
||
/// Called after each training epoch
|
||
///
|
||
/// # Arguments
|
||
///
|
||
/// * `trial_num` - Trial number
|
||
/// * `epoch` - Epoch number (0-indexed)
|
||
/// * `metrics` - Epoch metrics (train/val loss, timestamp)
|
||
///
|
||
/// # Returns
|
||
///
|
||
/// Decision on whether to continue, stop trial, or stop study
|
||
fn on_epoch_complete(
|
||
&mut self,
|
||
trial_num: usize,
|
||
epoch: usize,
|
||
metrics: &EpochMetrics,
|
||
) -> ObserverDecision;
|
||
|
||
/// Called when trial completes successfully
|
||
///
|
||
/// # Arguments
|
||
///
|
||
/// * `trial_num` - Trial number
|
||
/// * `final_loss` - Final validation loss
|
||
fn on_trial_complete(&mut self, trial_num: usize, final_loss: f64);
|
||
|
||
/// Called when trial fails with error
|
||
///
|
||
/// # Arguments
|
||
///
|
||
/// * `trial_num` - Trial number
|
||
/// * `error` - Error message
|
||
fn on_trial_failed(&mut self, trial_num: usize, error: &str);
|
||
}
|
||
|
||
/// Default early stopping observer implementation
|
||
///
|
||
/// Manages early stopping state for multiple trials and applies
|
||
/// configured stopping strategy. Implements `TrialObserver` trait
|
||
/// for integration with training loops.
|
||
///
|
||
/// ## State Management
|
||
///
|
||
/// - Maintains per-trial state in `HashMap<trial_id, State>`
|
||
/// - Tracks global best loss for cross-trial strategies
|
||
/// - Records complete epoch history for analysis
|
||
///
|
||
/// ## Memory Usage
|
||
///
|
||
/// - Fixed: ~128 bytes (observer overhead)
|
||
/// - Per-trial: ~48 + 48N bytes (N = epochs)
|
||
/// - Total for 30 trials × 100 epochs: ~144KB
|
||
///
|
||
/// ## Example
|
||
///
|
||
/// ```rust
|
||
/// use ml::hyperopt::early_stopping::{
|
||
/// EarlyStoppingConfig, EarlyStoppingObserver,
|
||
/// EarlyStoppingStrategy, TrialObserver, EpochMetrics,
|
||
/// };
|
||
///
|
||
/// let config = EarlyStoppingConfig {
|
||
/// patience_epochs: 10,
|
||
/// min_delta: 1e-4,
|
||
/// strategy: EarlyStoppingStrategy::Plateau,
|
||
/// ..Default::default()
|
||
/// };
|
||
///
|
||
/// let mut observer = EarlyStoppingObserver::new(config);
|
||
///
|
||
/// // In adapter's train_with_params()
|
||
/// observer.on_trial_start(0, "lr=0.001");
|
||
///
|
||
/// for epoch in 0..100 {
|
||
/// // ... training ...
|
||
/// let metrics = EpochMetrics {
|
||
/// epoch,
|
||
/// train_loss: 0.4,
|
||
/// val_loss: 0.5,
|
||
/// timestamp: epoch as f64,
|
||
/// };
|
||
///
|
||
/// match observer.on_epoch_complete(0, epoch, &metrics) {
|
||
/// ml::hyperopt::early_stopping::ObserverDecision::StopTrial => break,
|
||
/// _ => continue,
|
||
/// }
|
||
/// }
|
||
///
|
||
/// observer.on_trial_complete(0, 0.5);
|
||
/// ```
|
||
#[derive(Debug)]
|
||
pub struct EarlyStoppingObserver {
|
||
config: EarlyStoppingConfig,
|
||
state: HashMap<usize, EarlyStoppingState>,
|
||
baseline_val_loss: Option<f64>,
|
||
trial_best_losses: Vec<f64>,
|
||
}
|
||
|
||
impl EarlyStoppingObserver {
|
||
/// Create new early stopping observer
|
||
///
|
||
/// # Arguments
|
||
///
|
||
/// * `config` - Early stopping configuration
|
||
pub fn new(config: EarlyStoppingConfig) -> Self {
|
||
Self {
|
||
config,
|
||
state: HashMap::new(),
|
||
baseline_val_loss: None,
|
||
trial_best_losses: Vec::new(),
|
||
}
|
||
}
|
||
|
||
/// Get number of trials tracked
|
||
pub fn trial_count(&self) -> usize {
|
||
self.state.len()
|
||
}
|
||
|
||
/// Get or create state for trial
|
||
fn get_or_create_state(&mut self, trial_num: usize) -> &mut EarlyStoppingState {
|
||
self.state
|
||
.entry(trial_num)
|
||
.or_default()
|
||
}
|
||
|
||
/// Check if should stop based on strategy
|
||
fn should_stop_trial(&self, epoch: usize, val_loss: f64, patience_counter: usize) -> bool {
|
||
// Always respect min_epochs
|
||
if epoch < self.config.min_epochs {
|
||
return false;
|
||
}
|
||
|
||
// Apply strategy-specific logic
|
||
match &self.config.strategy {
|
||
EarlyStoppingStrategy::Plateau => patience_counter >= self.config.patience_epochs,
|
||
EarlyStoppingStrategy::MedianPruner { warmup_steps } => {
|
||
let strategy = MedianPrunerStrategy::new(*warmup_steps);
|
||
strategy.should_stop(epoch, val_loss, &self.trial_best_losses)
|
||
},
|
||
EarlyStoppingStrategy::PercentilePruner {
|
||
percentile,
|
||
warmup_steps,
|
||
} => {
|
||
let strategy = PercentilePrunerStrategy::new(*percentile, *warmup_steps);
|
||
strategy.should_stop(epoch, val_loss, &self.trial_best_losses)
|
||
},
|
||
EarlyStoppingStrategy::SuccessiveHalving { reduction_factor } => {
|
||
// Need at least reduction_factor completed trials for meaningful comparison
|
||
if self.trial_best_losses.len() < *reduction_factor {
|
||
return false;
|
||
}
|
||
|
||
// Sort completed trial losses ascending (NaN-safe)
|
||
let mut sorted = self.trial_best_losses.clone();
|
||
sorted.sort_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal));
|
||
|
||
// Keep top 1/eta fraction, prune the rest
|
||
let keep_count = (sorted.len() / reduction_factor).max(1);
|
||
let threshold = sorted[keep_count - 1];
|
||
|
||
// Prune if current trial is worse than the threshold
|
||
val_loss > threshold
|
||
},
|
||
EarlyStoppingStrategy::Hyperband {
|
||
max_resource,
|
||
reduction_factor,
|
||
} => {
|
||
// Hyperband: SHA pruning but ONLY at rung epochs
|
||
// Rung epochs: max_resource / η^k for k = 1, 2, ...
|
||
let eta = *reduction_factor;
|
||
let max_r = *max_resource;
|
||
|
||
// Check if this epoch is a rung epoch
|
||
let mut is_rung = false;
|
||
let mut rung_r = max_r / eta;
|
||
while rung_r >= 1 {
|
||
if epoch == rung_r {
|
||
is_rung = true;
|
||
break;
|
||
}
|
||
rung_r /= eta;
|
||
}
|
||
|
||
if !is_rung {
|
||
return false; // Only prune at rung epochs
|
||
}
|
||
|
||
// At rung epoch, apply SHA pruning logic
|
||
if self.trial_best_losses.len() < eta {
|
||
return false;
|
||
}
|
||
|
||
let mut sorted_losses = self.trial_best_losses.clone();
|
||
sorted_losses.sort_by(|a, b| {
|
||
a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal)
|
||
});
|
||
|
||
let keep_count = (sorted_losses.len() / eta).max(1);
|
||
let threshold = sorted_losses[keep_count - 1];
|
||
|
||
val_loss > threshold
|
||
},
|
||
}
|
||
}
|
||
}
|
||
|
||
impl TrialObserver for EarlyStoppingObserver {
|
||
fn on_trial_start(&mut self, trial_num: usize, params: &str) {
|
||
info!(
|
||
"Early Stopping: Trial {} started with params: {}",
|
||
trial_num, params
|
||
);
|
||
self.get_or_create_state(trial_num);
|
||
}
|
||
|
||
fn on_epoch_complete(
|
||
&mut self,
|
||
trial_num: usize,
|
||
epoch: usize,
|
||
metrics: &EpochMetrics,
|
||
) -> ObserverDecision {
|
||
let val_loss = metrics.val_loss;
|
||
let min_delta = self.config.min_delta;
|
||
let patience_epochs = self.config.patience_epochs;
|
||
|
||
// Get state
|
||
let state = self.get_or_create_state(trial_num);
|
||
|
||
// Record metrics
|
||
state.record_epoch_metrics(metrics.clone());
|
||
|
||
// Update best loss
|
||
let improved = state.update(val_loss, min_delta);
|
||
|
||
if improved {
|
||
info!(
|
||
"Trial {}, Epoch {}: val_loss improved to {:.6} (patience reset)",
|
||
trial_num, epoch, val_loss
|
||
);
|
||
}
|
||
|
||
// Check stopping criteria
|
||
// Extract patience_counter before checking (to avoid borrow issues)
|
||
let patience_counter = state.patience_counter;
|
||
|
||
// Drop mutable borrow of state before calling should_stop_trial
|
||
let should_stop = self.should_stop_trial(epoch, val_loss, patience_counter);
|
||
|
||
if should_stop {
|
||
// Get state again to mark as stopped
|
||
let state = self.get_or_create_state(trial_num);
|
||
state.mark_stopped(epoch);
|
||
info!(
|
||
"Trial {}, Epoch {}: Early stopping triggered (patience: {}/{})",
|
||
trial_num, epoch, patience_counter, patience_epochs
|
||
);
|
||
return ObserverDecision::StopTrial;
|
||
}
|
||
|
||
ObserverDecision::Continue
|
||
}
|
||
|
||
fn on_trial_complete(&mut self, trial_num: usize, final_loss: f64) {
|
||
info!(
|
||
"Early Stopping: Trial {} completed with final loss: {:.6}",
|
||
trial_num, final_loss
|
||
);
|
||
|
||
// Update global statistics
|
||
self.trial_best_losses.push(final_loss);
|
||
|
||
if self.baseline_val_loss.is_none() || final_loss < self.baseline_val_loss.unwrap_or(f64::INFINITY) {
|
||
self.baseline_val_loss = Some(final_loss);
|
||
info!("New baseline loss: {:.6}", final_loss);
|
||
}
|
||
}
|
||
|
||
fn on_trial_failed(&mut self, trial_num: usize, error: &str) {
|
||
warn!("Early Stopping: Trial {} failed: {}", trial_num, error);
|
||
}
|
||
}
|
||
|
||
#[cfg(test)]
|
||
mod tests {
|
||
use super::*;
|
||
|
||
#[test]
|
||
fn test_early_stopping_state_update_logic() {
|
||
let mut state = EarlyStoppingState::new();
|
||
|
||
// First update always improves
|
||
assert!(state.update(1.0, 0.01));
|
||
assert_eq!(state.best_val_loss, 1.0);
|
||
assert_eq!(state.patience_counter, 0);
|
||
|
||
// Significant improvement (> 0.01)
|
||
assert!(state.update(0.5, 0.01));
|
||
assert_eq!(state.best_val_loss, 0.5);
|
||
assert_eq!(state.patience_counter, 0);
|
||
|
||
// Marginal improvement (< 0.01)
|
||
assert!(!state.update(0.49, 0.01));
|
||
assert_eq!(state.best_val_loss, 0.5);
|
||
assert_eq!(state.patience_counter, 1);
|
||
|
||
// No improvement
|
||
assert!(!state.update(0.6, 0.01));
|
||
assert_eq!(state.best_val_loss, 0.5);
|
||
assert_eq!(state.patience_counter, 2);
|
||
}
|
||
|
||
#[test]
|
||
fn test_plateau_detection_basic() {
|
||
let strategy = PlateauDetectionStrategy::new(3, 0.01);
|
||
let mut state = EarlyStoppingState::new();
|
||
|
||
// Improving - should not stop
|
||
assert!(!strategy.should_stop(0, 1.0, &mut state, &[]));
|
||
assert!(!strategy.should_stop(1, 0.5, &mut state, &[]));
|
||
|
||
// Plateau - should stop after patience
|
||
assert!(!strategy.should_stop(2, 0.5, &mut state, &[]));
|
||
assert!(!strategy.should_stop(3, 0.5, &mut state, &[]));
|
||
assert!(strategy.should_stop(4, 0.5, &mut state, &[]));
|
||
}
|
||
|
||
#[test]
|
||
fn test_median_pruner_basic() {
|
||
let strategy = MedianPrunerStrategy::new(0);
|
||
let history = vec![0.3, 0.4, 0.5, 0.6, 0.7]; // median = 0.5
|
||
|
||
// Worse than median
|
||
assert!(strategy.should_stop(5, 0.8, &history));
|
||
|
||
// Better than median
|
||
assert!(!strategy.should_stop(5, 0.2, &history));
|
||
|
||
// Equal to median
|
||
assert!(!strategy.should_stop(5, 0.5, &history));
|
||
}
|
||
|
||
#[test]
|
||
fn test_percentile_pruner_basic() {
|
||
let strategy = PercentilePrunerStrategy::new(25.0, 0);
|
||
let history = vec![0.2, 0.4, 0.6, 0.8]; // 25th = 0.4
|
||
|
||
// Bottom 25%
|
||
assert!(strategy.should_stop(5, 0.9, &history));
|
||
|
||
// Above 25th percentile (0.3 < 0.4, so not pruned)
|
||
assert!(!strategy.should_stop(5, 0.3, &history));
|
||
}
|
||
|
||
#[test]
|
||
fn test_successive_halving_prunes_bottom_trials() {
|
||
// 9 completed trials with losses 0.1..=0.9, reduction_factor=3
|
||
// keep_count = max(1, 9/3) = 3
|
||
// sorted: [0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9]
|
||
// threshold = sorted[2] = 0.3
|
||
// New trial with val_loss=0.5 > 0.3 => should be pruned (StopTrial)
|
||
let config = EarlyStoppingConfig {
|
||
patience_epochs: 10,
|
||
min_delta: 1e-4,
|
||
min_epochs: 0,
|
||
strategy: EarlyStoppingStrategy::SuccessiveHalving {
|
||
reduction_factor: 3,
|
||
},
|
||
..Default::default()
|
||
};
|
||
let mut observer = EarlyStoppingObserver::new(config);
|
||
|
||
// Register 9 completed trials
|
||
for i in 0..9 {
|
||
observer.on_trial_start(i, &format!("trial_{}", i));
|
||
let loss = (i as f64 + 1.0) / 10.0; // 0.1, 0.2, ..., 0.9
|
||
observer.on_trial_complete(i, loss);
|
||
}
|
||
|
||
// Start a new trial (trial 9) with a bad loss
|
||
observer.on_trial_start(9, "trial_9");
|
||
let metrics = EpochMetrics {
|
||
epoch: 1,
|
||
train_loss: 0.5,
|
||
val_loss: 0.5, // worse than threshold 0.3
|
||
timestamp: 1.0,
|
||
};
|
||
let decision = observer.on_epoch_complete(9, 1, &metrics);
|
||
assert_eq!(
|
||
decision,
|
||
ObserverDecision::StopTrial,
|
||
"Trial with val_loss=0.5 should be pruned (threshold=0.3)"
|
||
);
|
||
}
|
||
|
||
#[test]
|
||
fn test_successive_halving_keeps_top_trials() {
|
||
// 9 completed trials with losses 0.1..=0.9, reduction_factor=3
|
||
// keep_count = max(1, 9/3) = 3
|
||
// sorted: [0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9]
|
||
// threshold = sorted[2] = 0.3
|
||
// New trial with val_loss=0.2 <= 0.3 => should continue
|
||
let config = EarlyStoppingConfig {
|
||
patience_epochs: 10,
|
||
min_delta: 1e-4,
|
||
min_epochs: 0,
|
||
strategy: EarlyStoppingStrategy::SuccessiveHalving {
|
||
reduction_factor: 3,
|
||
},
|
||
..Default::default()
|
||
};
|
||
let mut observer = EarlyStoppingObserver::new(config);
|
||
|
||
// Register 9 completed trials
|
||
for i in 0..9 {
|
||
observer.on_trial_start(i, &format!("trial_{}", i));
|
||
let loss = (i as f64 + 1.0) / 10.0; // 0.1, 0.2, ..., 0.9
|
||
observer.on_trial_complete(i, loss);
|
||
}
|
||
|
||
// Start a new trial (trial 9) with a good loss
|
||
observer.on_trial_start(9, "trial_9");
|
||
let metrics = EpochMetrics {
|
||
epoch: 1,
|
||
train_loss: 0.2,
|
||
val_loss: 0.2, // within top 1/3 (threshold=0.3)
|
||
timestamp: 1.0,
|
||
};
|
||
let decision = observer.on_epoch_complete(9, 1, &metrics);
|
||
assert_eq!(
|
||
decision,
|
||
ObserverDecision::Continue,
|
||
"Trial with val_loss=0.2 should continue (threshold=0.3)"
|
||
);
|
||
}
|
||
|
||
#[test]
|
||
fn test_hyperband_prunes_at_rung_epochs() {
|
||
// max_resource=81, reduction_factor=3
|
||
// Rung epochs: 81/3=27, 81/9=9, 81/27=3, 81/81=1
|
||
let config = EarlyStoppingConfig {
|
||
patience_epochs: 100,
|
||
min_delta: 1e-4,
|
||
min_epochs: 0,
|
||
strategy: EarlyStoppingStrategy::Hyperband {
|
||
max_resource: 81,
|
||
reduction_factor: 3,
|
||
},
|
||
..Default::default()
|
||
};
|
||
let mut observer = EarlyStoppingObserver::new(config);
|
||
|
||
// Seed with 9 completed trials
|
||
for trial in 0..9 {
|
||
observer.on_trial_start(trial, &format!("trial_{}", trial));
|
||
observer.on_trial_complete(trial, (trial as f64 + 1.0) * 0.1);
|
||
}
|
||
|
||
// Trial 9 with bad loss at rung epoch 27 (81/3=27)
|
||
observer.on_trial_start(9, "trial_9");
|
||
let metrics = EpochMetrics {
|
||
epoch: 27,
|
||
train_loss: 0.9,
|
||
val_loss: 0.95,
|
||
timestamp: 27.0,
|
||
};
|
||
let decision = observer.on_epoch_complete(9, 27, &metrics);
|
||
assert_eq!(
|
||
decision,
|
||
ObserverDecision::StopTrial,
|
||
"Hyperband should prune at rung epoch 27 (val_loss=0.95 > threshold)"
|
||
);
|
||
}
|
||
|
||
#[test]
|
||
fn test_hyperband_does_not_prune_between_rungs() {
|
||
// max_resource=81, reduction_factor=3
|
||
// Rung epochs: 27, 9, 3, 1 -- epoch 15 is NOT a rung
|
||
let config = EarlyStoppingConfig {
|
||
patience_epochs: 100,
|
||
min_delta: 1e-4,
|
||
min_epochs: 0,
|
||
strategy: EarlyStoppingStrategy::Hyperband {
|
||
max_resource: 81,
|
||
reduction_factor: 3,
|
||
},
|
||
..Default::default()
|
||
};
|
||
let mut observer = EarlyStoppingObserver::new(config);
|
||
|
||
// Seed with 9 completed trials
|
||
for trial in 0..9 {
|
||
observer.on_trial_start(trial, &format!("trial_{}", trial));
|
||
observer.on_trial_complete(trial, (trial as f64 + 1.0) * 0.1);
|
||
}
|
||
|
||
// Trial 9 with bad loss but NOT at a rung epoch (epoch 15)
|
||
observer.on_trial_start(9, "trial_9");
|
||
let metrics = EpochMetrics {
|
||
epoch: 15,
|
||
train_loss: 0.9,
|
||
val_loss: 0.95,
|
||
timestamp: 15.0,
|
||
};
|
||
let decision = observer.on_epoch_complete(9, 15, &metrics);
|
||
assert_eq!(
|
||
decision,
|
||
ObserverDecision::Continue,
|
||
"Hyperband should NOT prune between rung epochs (epoch 15 is not a rung)"
|
||
);
|
||
}
|
||
}
|