feat: add ThresholdPruner for fixed-bound trial pruning
This commit is contained in:
+2
-2
@@ -201,7 +201,7 @@ pub use param::ParamValue;
|
||||
pub use parameter::{
|
||||
BoolParam, Categorical, CategoricalParam, EnumParam, FloatParam, IntParam, ParamId, Parameter,
|
||||
};
|
||||
pub use pruner::{NopPruner, Pruner};
|
||||
pub use pruner::{NopPruner, Pruner, ThresholdPruner};
|
||||
pub use sampler::CompletedTrial;
|
||||
pub use sampler::grid::GridSearchSampler;
|
||||
pub use sampler::random::RandomSampler;
|
||||
@@ -224,7 +224,7 @@ pub mod prelude {
|
||||
pub use crate::parameter::{
|
||||
BoolParam, Categorical, CategoricalParam, EnumParam, FloatParam, IntParam, Parameter,
|
||||
};
|
||||
pub use crate::pruner::{NopPruner, Pruner};
|
||||
pub use crate::pruner::{NopPruner, Pruner, ThresholdPruner};
|
||||
pub use crate::sampler::CompletedTrial;
|
||||
pub use crate::sampler::grid::GridSearchSampler;
|
||||
pub use crate::sampler::random::RandomSampler;
|
||||
|
||||
@@ -5,8 +5,10 @@
|
||||
//! discarding unpromising trials before they complete, saving compute.
|
||||
|
||||
mod nop;
|
||||
mod threshold;
|
||||
|
||||
pub use nop::NopPruner;
|
||||
pub use threshold::ThresholdPruner;
|
||||
|
||||
use crate::sampler::CompletedTrial;
|
||||
|
||||
|
||||
@@ -0,0 +1,83 @@
|
||||
use super::Pruner;
|
||||
use crate::sampler::CompletedTrial;
|
||||
|
||||
/// Prune trials whose intermediate values exceed fixed thresholds.
|
||||
///
|
||||
/// Useful for cutting off trials that are clearly diverging or stuck
|
||||
/// at bad values early in training.
|
||||
///
|
||||
/// # Examples
|
||||
///
|
||||
/// ```
|
||||
/// use optimizer::pruner::ThresholdPruner;
|
||||
///
|
||||
/// // Prune if the intermediate value exceeds 100.0 or falls below 0.0
|
||||
/// let pruner = ThresholdPruner::new().upper(100.0).lower(0.0);
|
||||
/// ```
|
||||
pub struct ThresholdPruner {
|
||||
/// Prune if intermediate value is greater than this. `None` = no upper bound.
|
||||
upper: Option<f64>,
|
||||
/// Prune if intermediate value is less than this. `None` = no lower bound.
|
||||
lower: Option<f64>,
|
||||
}
|
||||
|
||||
impl ThresholdPruner {
|
||||
/// Create a new `ThresholdPruner` with no thresholds set.
|
||||
///
|
||||
/// By default, no pruning occurs. Use [`upper`](Self::upper) and
|
||||
/// [`lower`](Self::lower) to set bounds.
|
||||
#[must_use]
|
||||
pub fn new() -> Self {
|
||||
Self {
|
||||
upper: None,
|
||||
lower: None,
|
||||
}
|
||||
}
|
||||
|
||||
/// Set the upper threshold. Trials with intermediate values above this
|
||||
/// will be pruned.
|
||||
#[must_use]
|
||||
pub fn upper(mut self, threshold: f64) -> Self {
|
||||
self.upper = Some(threshold);
|
||||
self
|
||||
}
|
||||
|
||||
/// Set the lower threshold. Trials with intermediate values below this
|
||||
/// will be pruned.
|
||||
#[must_use]
|
||||
pub fn lower(mut self, threshold: f64) -> Self {
|
||||
self.lower = Some(threshold);
|
||||
self
|
||||
}
|
||||
}
|
||||
|
||||
impl Default for ThresholdPruner {
|
||||
fn default() -> Self {
|
||||
Self::new()
|
||||
}
|
||||
}
|
||||
|
||||
impl Pruner for ThresholdPruner {
|
||||
fn should_prune(
|
||||
&self,
|
||||
_trial_id: u64,
|
||||
_step: u64,
|
||||
intermediate_values: &[(u64, f64)],
|
||||
_completed_trials: &[CompletedTrial],
|
||||
) -> bool {
|
||||
let Some(&(_, latest_value)) = intermediate_values.last() else {
|
||||
return false;
|
||||
};
|
||||
if let Some(upper) = self.upper
|
||||
&& latest_value > upper
|
||||
{
|
||||
return true;
|
||||
}
|
||||
if let Some(lower) = self.lower
|
||||
&& latest_value < lower
|
||||
{
|
||||
return true;
|
||||
}
|
||||
false
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,64 @@
|
||||
use optimizer::pruner::{Pruner, ThresholdPruner};
|
||||
|
||||
#[test]
|
||||
fn prune_when_value_exceeds_upper_threshold() {
|
||||
let pruner = ThresholdPruner::new().upper(10.0);
|
||||
let values = vec![(0, 5.0), (1, 8.0), (2, 11.0)];
|
||||
assert!(pruner.should_prune(0, 2, &values, &[]));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn prune_when_value_falls_below_lower_threshold() {
|
||||
let pruner = ThresholdPruner::new().lower(0.0);
|
||||
let values = vec![(0, 5.0), (1, 2.0), (2, -1.0)];
|
||||
assert!(pruner.should_prune(0, 2, &values, &[]));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn no_prune_when_value_within_bounds() {
|
||||
let pruner = ThresholdPruner::new().upper(10.0).lower(0.0);
|
||||
let values = vec![(0, 3.0), (1, 5.0), (2, 7.0)];
|
||||
assert!(!pruner.should_prune(0, 2, &values, &[]));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn no_prune_when_no_intermediate_values() {
|
||||
let pruner = ThresholdPruner::new().upper(10.0).lower(0.0);
|
||||
assert!(!pruner.should_prune(0, 0, &[], &[]));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn works_with_only_upper_set() {
|
||||
let pruner = ThresholdPruner::new().upper(5.0);
|
||||
let below = vec![(0, 3.0)];
|
||||
let above = vec![(0, 6.0)];
|
||||
assert!(!pruner.should_prune(0, 0, &below, &[]));
|
||||
assert!(pruner.should_prune(0, 0, &above, &[]));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn works_with_only_lower_set() {
|
||||
let pruner = ThresholdPruner::new().lower(2.0);
|
||||
let above = vec![(0, 5.0)];
|
||||
let below = vec![(0, 1.0)];
|
||||
assert!(!pruner.should_prune(0, 0, &above, &[]));
|
||||
assert!(pruner.should_prune(0, 0, &below, &[]));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn works_with_both_thresholds_set() {
|
||||
let pruner = ThresholdPruner::new().upper(10.0).lower(0.0);
|
||||
|
||||
// Within bounds
|
||||
assert!(!pruner.should_prune(0, 0, &[(0, 5.0)], &[]));
|
||||
|
||||
// Exceeds upper
|
||||
assert!(pruner.should_prune(0, 0, &[(0, 15.0)], &[]));
|
||||
|
||||
// Below lower
|
||||
assert!(pruner.should_prune(0, 0, &[(0, -3.0)], &[]));
|
||||
|
||||
// At exact boundary (not pruned — strictly greater/less)
|
||||
assert!(!pruner.should_prune(0, 0, &[(0, 10.0)], &[]));
|
||||
assert!(!pruner.should_prune(0, 0, &[(0, 0.0)], &[]));
|
||||
}
|
||||
Reference in New Issue
Block a user