feat: add ThresholdPruner for fixed-bound trial pruning
This commit is contained in:
@@ -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
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user