Manuel Raimann
68 lines
2.4 KiB
Rust
68 lines
2.4 KiB
Rust
//! Trial pruning — stop unpromising trials early with `MedianPruner`.
|
|
//!
|
|
//! When your objective involves an iterative loop (e.g. training epochs),
|
|
//! the pruner compares intermediate values across trials and kills the
|
|
//! ones that fall below the median — saving compute on bad configurations.
|
|
//!
|
|
//! Run with: `cargo run --example pruning`
|
|
|
|
use optimizer::prelude::*;
|
|
|
|
fn main() -> optimizer::Result<()> {
|
|
// MedianPruner prunes trials whose intermediate value falls below the
|
|
// median of previously completed trials at the same step.
|
|
let study: Study<f64> = Study::builder()
|
|
.minimize()
|
|
.sampler(RandomSampler::with_seed(42))
|
|
.pruner(
|
|
MedianPruner::new(Direction::Minimize)
|
|
.n_warmup_steps(3) // run at least 3 epochs before pruning
|
|
.n_min_trials(3), // need 3 completed trials before pruning kicks in
|
|
)
|
|
.build();
|
|
|
|
let lr = FloatParam::new(1e-4, 1.0).name("learning_rate");
|
|
let momentum = FloatParam::new(0.0, 0.99).name("momentum");
|
|
|
|
let n_epochs: u64 = 20;
|
|
|
|
study.optimize(30, |trial: &mut optimizer::Trial| {
|
|
let lr_val = lr.suggest(trial)?;
|
|
let mom = momentum.suggest(trial)?;
|
|
|
|
// Simulated training loop — good hyperparameters converge to low loss,
|
|
// bad ones plateau high, giving the pruner something to cut.
|
|
let mut loss = 1.0;
|
|
for epoch in 0..n_epochs {
|
|
let lr_penalty = (lr_val.log10() - 0.01_f64.log10()).powi(2);
|
|
let mom_penalty = (mom - 0.8).powi(2);
|
|
let base_loss = 0.02 + 0.05 * lr_penalty + 1.5 * mom_penalty;
|
|
let progress = (epoch as f64 + 1.0) / n_epochs as f64;
|
|
loss = base_loss + (1.0 - base_loss) * (-3.5 * progress).exp();
|
|
|
|
// Report intermediate value so the pruner can evaluate this trial.
|
|
trial.report(epoch, loss);
|
|
|
|
// Check whether the pruner recommends stopping early.
|
|
if trial.should_prune() {
|
|
Err(TrialPruned)?;
|
|
}
|
|
}
|
|
|
|
Ok::<_, Error>(loss)
|
|
})?;
|
|
|
|
// --- Results ---
|
|
let best = study.best_trial()?;
|
|
println!(
|
|
"Recorded {} trials ({} pruned)",
|
|
study.n_trials(),
|
|
study.n_pruned_trials()
|
|
);
|
|
println!("Best trial #{}: loss = {:.6}", best.id, best.value);
|
|
println!(" learning_rate = {:.6}", best.get(&lr).unwrap());
|
|
println!(" momentum = {:.4}", best.get(&momentum).unwrap());
|
|
|
|
Ok(())
|
|
}
|