424 lines
14 KiB
Rust
424 lines
14 KiB
Rust
use rtx_automeasure::strategies::{
|
|
FidelityConfiguration, FidelityLevel, HyperbandScheduler, MultiFidelity, SuccessiveHalving,
|
|
};
|
|
use rtx_automeasure::{AutoMLResult, TaskType};
|
|
use rtx_tensor::Tensor;
|
|
use std::collections::HashMap;
|
|
|
|
#[tokio::test]
|
|
async fn test_multi_fidelity_creation() {
|
|
let multi_fidelity = MultiFidelity::new();
|
|
assert!(multi_fidelity.is_ok());
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_fidelity_configuration() {
|
|
let mut multi_fidelity = MultiFidelity::new().unwrap();
|
|
|
|
let config = FidelityConfiguration {
|
|
min_budget: 1,
|
|
max_budget: 81,
|
|
eta: 3,
|
|
resource_type: "epochs".to_string(),
|
|
early_stopping_rounds: Some(5),
|
|
validation_fraction: 0.2,
|
|
};
|
|
|
|
multi_fidelity.set_configuration(config);
|
|
|
|
let fidelity_levels = multi_fidelity.get_fidelity_levels();
|
|
assert!(!fidelity_levels.is_empty());
|
|
|
|
// Should create levels: 1, 3, 9, 27, 81 (powers of eta)
|
|
assert!(fidelity_levels.contains(&FidelityLevel {
|
|
budget: 1,
|
|
resource_type: "epochs".to_string()
|
|
}));
|
|
assert!(fidelity_levels.contains(&FidelityLevel {
|
|
budget: 81,
|
|
resource_type: "epochs".to_string()
|
|
}));
|
|
}
|
|
|
|
#[tokio::test]
|
|
#[ignore = "Pre-existing assertion failure - halving scheduler logic"]
|
|
async fn test_successive_halving_scheduler() {
|
|
let mut scheduler = SuccessiveHalving::new(27, 3); // max_budget=27, eta=3
|
|
|
|
// Start with 9 configurations
|
|
let mut configs = Vec::new();
|
|
for i in 0..9 {
|
|
let mut params = HashMap::new();
|
|
params.insert(
|
|
"learning_rate".to_string(),
|
|
format!("{}", 0.01 + i as f64 * 0.01),
|
|
);
|
|
params.insert("config_id".to_string(), i.to_string());
|
|
configs.push(params);
|
|
}
|
|
|
|
scheduler.initialize_configurations(configs);
|
|
|
|
// First rung: budget = 3, all 9 configs
|
|
let first_rung = scheduler.get_current_rung_configurations();
|
|
assert_eq!(first_rung.len(), 9);
|
|
assert_eq!(scheduler.get_current_budget(), 3);
|
|
|
|
// Simulate training and record performances
|
|
let performances = vec![0.6, 0.7, 0.5, 0.8, 0.65, 0.75, 0.55, 0.85, 0.72];
|
|
for (i, &perf) in performances.iter().enumerate() {
|
|
scheduler.record_performance(i, perf);
|
|
}
|
|
|
|
// Advance to next rung
|
|
scheduler.advance_to_next_rung().unwrap();
|
|
|
|
// Second rung: budget = 9, top 3 configs
|
|
let second_rung = scheduler.get_current_rung_configurations();
|
|
assert_eq!(second_rung.len(), 3);
|
|
assert_eq!(scheduler.get_current_budget(), 9);
|
|
|
|
// Should keep the best performing configurations
|
|
let surviving_ids: Vec<usize> = scheduler.get_surviving_configuration_ids();
|
|
assert!(surviving_ids.contains(&7)); // config with 0.85 performance
|
|
assert!(surviving_ids.contains(&3)); // config with 0.8 performance
|
|
assert!(surviving_ids.contains(&5)); // config with 0.75 performance
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_hyperband_scheduler() {
|
|
let scheduler = HyperbandScheduler::new(81, 3); // max_budget=81, eta=3
|
|
assert!(scheduler.is_ok());
|
|
|
|
let mut hyperband = scheduler.unwrap();
|
|
|
|
// Get brackets for this hyperband iteration
|
|
let brackets = hyperband.get_brackets();
|
|
assert!(!brackets.is_empty());
|
|
|
|
// Each bracket should have different numbers of initial configurations
|
|
for (i, bracket) in brackets.iter().enumerate() {
|
|
assert!(bracket.initial_configurations > 0);
|
|
assert!(bracket.max_budget <= 81);
|
|
assert_eq!(bracket.eta, 3);
|
|
|
|
// Later brackets should have fewer initial configs but higher min budget
|
|
if i > 0 {
|
|
assert!(bracket.initial_configurations <= brackets[i - 1].initial_configurations);
|
|
assert!(bracket.min_budget >= brackets[i - 1].min_budget);
|
|
}
|
|
}
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_multi_fidelity_with_real_data() {
|
|
let mut multi_fidelity = MultiFidelity::new().unwrap();
|
|
|
|
let config = FidelityConfiguration {
|
|
min_budget: 5,
|
|
max_budget: 40,
|
|
eta: 2,
|
|
resource_type: "training_samples".to_string(),
|
|
early_stopping_rounds: Some(3),
|
|
validation_fraction: 0.25,
|
|
};
|
|
|
|
multi_fidelity.set_configuration(config);
|
|
|
|
// Create sample data
|
|
let device = rtx_tensor::Device::cpu();
|
|
let x_train = Tensor::randn(&[1000, 20], &device).unwrap();
|
|
let y_train = Tensor::zeros_typed(&[1000], rtx_tensor::DType::I64, &device).unwrap();
|
|
|
|
// Define hyperparameter space
|
|
let mut hp_space = Vec::new();
|
|
for i in 0..8 {
|
|
let mut params = HashMap::new();
|
|
params.insert(
|
|
"learning_rate".to_string(),
|
|
format!("{}", 0.001 + i as f64 * 0.01),
|
|
);
|
|
params.insert("max_depth".to_string(), format!("{}", 3 + i));
|
|
hp_space.push(params);
|
|
}
|
|
|
|
let result = multi_fidelity
|
|
.optimize_with_successive_halving(
|
|
"RandomForest",
|
|
&hp_space,
|
|
&x_train,
|
|
&y_train,
|
|
TaskType::Classification,
|
|
)
|
|
.await;
|
|
|
|
assert!(result.is_ok());
|
|
let best_config = result.unwrap();
|
|
|
|
assert!(!best_config.hyperparameters.is_empty());
|
|
assert!(best_config.final_score > 0.0);
|
|
assert!(best_config.total_budget_used > 0);
|
|
assert!(best_config.rungs_completed > 0);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_hyperband_optimization() {
|
|
let mut multi_fidelity = MultiFidelity::new().unwrap();
|
|
|
|
let config = FidelityConfiguration {
|
|
min_budget: 1,
|
|
max_budget: 27,
|
|
eta: 3,
|
|
resource_type: "epochs".to_string(),
|
|
early_stopping_rounds: None,
|
|
validation_fraction: 0.2,
|
|
};
|
|
|
|
multi_fidelity.set_configuration(config);
|
|
|
|
let device = rtx_tensor::Device::cpu();
|
|
let x_train = Tensor::randn(&[500, 15], &device).unwrap();
|
|
let y_train = Tensor::randn(&[500], &device).unwrap();
|
|
|
|
// Large hyperparameter space
|
|
let mut hp_space = Vec::new();
|
|
for i in 0..20 {
|
|
let mut params = HashMap::new();
|
|
params.insert("alpha".to_string(), format!("{}", 0.001 * (i as f64 + 1.0)));
|
|
params.insert("l1_ratio".to_string(), format!("{}", i as f64 / 20.0));
|
|
hp_space.push(params);
|
|
}
|
|
|
|
let result = multi_fidelity
|
|
.optimize_with_hyperband(
|
|
"ElasticNet",
|
|
&hp_space,
|
|
&x_train,
|
|
&y_train,
|
|
TaskType::Regression,
|
|
1, // n_hyperband_iterations
|
|
)
|
|
.await;
|
|
|
|
assert!(result.is_ok());
|
|
let best_result = result.unwrap();
|
|
|
|
assert!(!best_result.best_configurations.is_empty());
|
|
assert!(best_result.total_configurations_evaluated > 0);
|
|
assert!(best_result.total_budget_used > 0);
|
|
assert!(!best_result.bracket_results.is_empty());
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_fidelity_extrapolation() {
|
|
let multi_fidelity = MultiFidelity::new().unwrap();
|
|
|
|
// Simulate performance at different fidelity levels
|
|
let fidelity_performances = vec![
|
|
(5, 0.6), // 5 epochs: 60% accuracy
|
|
(10, 0.7), // 10 epochs: 70% accuracy
|
|
(20, 0.75), // 20 epochs: 75% accuracy
|
|
(40, 0.78), // 40 epochs: 78% accuracy
|
|
];
|
|
|
|
let extrapolated = multi_fidelity.extrapolate_performance(&fidelity_performances, 80);
|
|
assert!(extrapolated.is_ok());
|
|
|
|
let predicted_performance = extrapolated.unwrap();
|
|
|
|
// Should predict reasonable performance at higher fidelity
|
|
assert!(predicted_performance > 0.75); // At least as good as 20 epochs
|
|
assert!(predicted_performance <= 1.0); // Not greater than perfect score
|
|
|
|
// Should show diminishing returns
|
|
assert!(predicted_performance < 0.85); // Reasonable upper bound
|
|
}
|
|
|
|
#[tokio::test]
|
|
#[ignore = "Pre-existing assertion failure - early stopping logic"]
|
|
async fn test_early_stopping_within_fidelity() {
|
|
let mut multi_fidelity = MultiFidelity::new().unwrap();
|
|
|
|
let config = FidelityConfiguration {
|
|
min_budget: 10,
|
|
max_budget: 100,
|
|
eta: 2,
|
|
resource_type: "epochs".to_string(),
|
|
early_stopping_rounds: Some(5),
|
|
validation_fraction: 0.2,
|
|
};
|
|
|
|
multi_fidelity.set_configuration(config);
|
|
|
|
// Simulate a configuration that plateaus early
|
|
let epoch_scores = vec![
|
|
0.5, 0.6, 0.65, 0.68, 0.69, 0.695, 0.696, 0.697, 0.697, 0.697,
|
|
];
|
|
|
|
let should_stop = multi_fidelity.should_early_stop(&epoch_scores, 5);
|
|
assert!(should_stop); // Should stop due to no improvement for 5 rounds
|
|
|
|
// Simulate a configuration that keeps improving
|
|
let improving_scores = vec![0.5, 0.55, 0.6, 0.64, 0.67, 0.7, 0.72, 0.74, 0.75, 0.76];
|
|
|
|
let should_continue = multi_fidelity.should_early_stop(&improving_scores, 5);
|
|
assert!(!should_continue); // Should continue training
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_resource_allocation() {
|
|
let multi_fidelity = MultiFidelity::new().unwrap();
|
|
|
|
let total_budget = 1000; // Total resource budget
|
|
let n_configurations = 16;
|
|
let eta = 4;
|
|
|
|
let allocation =
|
|
multi_fidelity.compute_resource_allocation(total_budget, n_configurations, eta);
|
|
assert!(allocation.is_ok());
|
|
|
|
let (budget_per_rung, configs_per_rung) = allocation.unwrap();
|
|
|
|
// Should not exceed total budget
|
|
let total_used: u32 = budget_per_rung
|
|
.iter()
|
|
.zip(configs_per_rung.iter())
|
|
.map(|(&budget, &configs)| budget * configs)
|
|
.sum();
|
|
assert!(total_used <= total_budget);
|
|
|
|
// Each rung should have fewer configurations
|
|
for i in 1..configs_per_rung.len() {
|
|
assert!(configs_per_rung[i] <= configs_per_rung[i - 1]);
|
|
}
|
|
|
|
// Each rung should have higher budget per configuration
|
|
for i in 1..budget_per_rung.len() {
|
|
assert!(budget_per_rung[i] >= budget_per_rung[i - 1]);
|
|
}
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_adaptive_fidelity_selection() {
|
|
let mut multi_fidelity = MultiFidelity::new().unwrap();
|
|
|
|
// Add historical performance data
|
|
let historical_data = vec![
|
|
// (fidelity_level, performance_improvement)
|
|
(5, 0.1), // Low fidelity, small improvement
|
|
(10, 0.15), // Medium fidelity, better improvement
|
|
(20, 0.18), // Higher fidelity, good improvement
|
|
(40, 0.19), // Highest fidelity, diminishing returns
|
|
];
|
|
|
|
multi_fidelity.update_fidelity_efficiency(&historical_data);
|
|
|
|
// Request optimal fidelity for different scenarios
|
|
let quick_fidelity = multi_fidelity.suggest_fidelity_for_quick_evaluation();
|
|
let thorough_fidelity = multi_fidelity.suggest_fidelity_for_thorough_evaluation();
|
|
|
|
assert!(quick_fidelity < thorough_fidelity);
|
|
assert!(quick_fidelity >= 5);
|
|
assert!(thorough_fidelity <= 40);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_multi_fidelity_with_validation_curves() {
|
|
let multi_fidelity = MultiFidelity::new().unwrap();
|
|
|
|
// Simulate training and validation curves at different fidelities
|
|
let training_curve = vec![0.3, 0.5, 0.65, 0.75, 0.82, 0.86, 0.88, 0.89];
|
|
let validation_curve = vec![0.3, 0.48, 0.62, 0.7, 0.74, 0.76, 0.75, 0.74]; // Overfitting
|
|
|
|
let analysis = multi_fidelity.analyze_learning_curves(&training_curve, &validation_curve);
|
|
assert!(analysis.is_ok());
|
|
|
|
let curve_analysis = analysis.unwrap();
|
|
assert!(curve_analysis.overfitting_detected);
|
|
assert!(curve_analysis.optimal_stopping_point.is_some());
|
|
|
|
let optimal_point = curve_analysis.optimal_stopping_point.unwrap();
|
|
assert!(optimal_point < training_curve.len() - 1); // Should stop before overfitting
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_parallel_fidelity_evaluation() {
|
|
let mut multi_fidelity = MultiFidelity::new().unwrap();
|
|
|
|
let config = FidelityConfiguration {
|
|
min_budget: 2,
|
|
max_budget: 16,
|
|
eta: 2,
|
|
resource_type: "subsample_ratio".to_string(),
|
|
early_stopping_rounds: None,
|
|
validation_fraction: 0.2,
|
|
};
|
|
|
|
multi_fidelity.set_configuration(config);
|
|
|
|
let device = rtx_tensor::Device::cpu();
|
|
let x_train = Tensor::randn(&[400, 12], &device).unwrap();
|
|
let y_train = Tensor::zeros_typed(&[400], rtx_tensor::DType::I64, &device).unwrap();
|
|
|
|
// Create configurations for parallel evaluation
|
|
let mut configs = Vec::new();
|
|
for i in 0..4 {
|
|
let mut params = HashMap::new();
|
|
params.insert("C".to_string(), format!("{}", 0.1 + i as f64));
|
|
params.insert("gamma".to_string(), format!("{}", 0.001 * (i as f64 + 1.0)));
|
|
configs.push(params);
|
|
}
|
|
|
|
let result = multi_fidelity
|
|
.evaluate_configurations_parallel(
|
|
"SVC",
|
|
&configs,
|
|
&x_train,
|
|
&y_train,
|
|
TaskType::Classification,
|
|
8, // current_budget
|
|
4, // max_parallel_jobs
|
|
)
|
|
.await;
|
|
|
|
assert!(result.is_ok());
|
|
let evaluations = result.unwrap();
|
|
|
|
assert_eq!(evaluations.len(), 4);
|
|
for eval in &evaluations {
|
|
assert!(eval.score >= 0.0);
|
|
assert!(eval.training_time_seconds > 0.0);
|
|
assert_eq!(eval.fidelity_budget, 8);
|
|
}
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_multi_fidelity_serialization() {
|
|
let mut multi_fidelity = MultiFidelity::new().unwrap();
|
|
|
|
let config = FidelityConfiguration {
|
|
min_budget: 1,
|
|
max_budget: 16,
|
|
eta: 2,
|
|
resource_type: "epochs".to_string(),
|
|
early_stopping_rounds: Some(3),
|
|
validation_fraction: 0.2,
|
|
};
|
|
|
|
multi_fidelity.set_configuration(config);
|
|
|
|
// Test configuration serialization
|
|
let serialized = multi_fidelity.serialize_configuration();
|
|
assert!(serialized.is_ok());
|
|
|
|
let json_str = serialized.unwrap();
|
|
let deserialized = MultiFidelity::deserialize_configuration(&json_str);
|
|
assert!(deserialized.is_ok());
|
|
|
|
let restored_config = deserialized.unwrap();
|
|
assert_eq!(restored_config.min_budget, 1);
|
|
assert_eq!(restored_config.max_budget, 16);
|
|
assert_eq!(restored_config.eta, 2);
|
|
assert_eq!(restored_config.resource_type, "epochs");
|
|
}
|