feat: add intermediate values and pruner integration to Trial
- Add report(step, value) and should_prune() methods to Trial - Store intermediate_values in Trial, copied to CompletedTrial on completion - Pass pruner through trial factory so trials can consult it - Rebuild trial factory when pruner is changed via set_pruner()
This commit is contained in:
@@ -27,6 +27,8 @@ pub struct CompletedTrial<V = f64> {
|
|||||||
pub param_labels: HashMap<ParamId, String>,
|
pub param_labels: HashMap<ParamId, String>,
|
||||||
/// The objective value returned by the objective function.
|
/// The objective value returned by the objective function.
|
||||||
pub value: V,
|
pub value: V,
|
||||||
|
/// Intermediate objective values reported during the trial.
|
||||||
|
pub intermediate_values: Vec<(u64, f64)>,
|
||||||
}
|
}
|
||||||
|
|
||||||
impl<V> CompletedTrial<V> {
|
impl<V> CompletedTrial<V> {
|
||||||
@@ -44,6 +46,26 @@ impl<V> CompletedTrial<V> {
|
|||||||
distributions,
|
distributions,
|
||||||
param_labels,
|
param_labels,
|
||||||
value,
|
value,
|
||||||
|
intermediate_values: Vec::new(),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Creates a new completed trial with intermediate values.
|
||||||
|
pub fn with_intermediate_values(
|
||||||
|
id: u64,
|
||||||
|
params: HashMap<ParamId, ParamValue>,
|
||||||
|
distributions: HashMap<ParamId, Distribution>,
|
||||||
|
param_labels: HashMap<ParamId, String>,
|
||||||
|
value: V,
|
||||||
|
intermediate_values: Vec<(u64, f64)>,
|
||||||
|
) -> Self {
|
||||||
|
Self {
|
||||||
|
id,
|
||||||
|
params,
|
||||||
|
distributions,
|
||||||
|
param_labels,
|
||||||
|
value,
|
||||||
|
intermediate_values,
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|||||||
+25
-8
@@ -105,14 +105,16 @@ where
|
|||||||
let sampler: Arc<dyn Sampler> = Arc::new(sampler);
|
let sampler: Arc<dyn Sampler> = Arc::new(sampler);
|
||||||
let completed_trials = Arc::new(RwLock::new(Vec::new()));
|
let completed_trials = Arc::new(RwLock::new(Vec::new()));
|
||||||
|
|
||||||
|
let pruner: Arc<dyn Pruner> = Arc::new(NopPruner);
|
||||||
|
|
||||||
// For Study<f64>, set up a trial factory that provides sampler integration.
|
// For Study<f64>, set up a trial factory that provides sampler integration.
|
||||||
// This uses Any downcasting to check at runtime whether V = f64.
|
// This uses Any downcasting to check at runtime whether V = f64.
|
||||||
let trial_factory = Self::make_trial_factory(&sampler, &completed_trials);
|
let trial_factory = Self::make_trial_factory(&sampler, &completed_trials, &pruner);
|
||||||
|
|
||||||
Self {
|
Self {
|
||||||
direction,
|
direction,
|
||||||
sampler,
|
sampler,
|
||||||
pruner: Arc::new(NopPruner),
|
pruner,
|
||||||
completed_trials,
|
completed_trials,
|
||||||
next_trial_id: AtomicU64::new(0),
|
next_trial_id: AtomicU64::new(0),
|
||||||
trial_factory,
|
trial_factory,
|
||||||
@@ -123,6 +125,7 @@ where
|
|||||||
fn make_trial_factory(
|
fn make_trial_factory(
|
||||||
sampler: &Arc<dyn Sampler>,
|
sampler: &Arc<dyn Sampler>,
|
||||||
completed_trials: &Arc<RwLock<Vec<CompletedTrial<V>>>>,
|
completed_trials: &Arc<RwLock<Vec<CompletedTrial<V>>>>,
|
||||||
|
pruner: &Arc<dyn Pruner>,
|
||||||
) -> Option<Arc<dyn Fn(u64) -> Trial + Send + Sync>>
|
) -> Option<Arc<dyn Fn(u64) -> Trial + Send + Sync>>
|
||||||
where
|
where
|
||||||
V: 'static,
|
V: 'static,
|
||||||
@@ -135,8 +138,14 @@ where
|
|||||||
f64_trials.map(|trials| {
|
f64_trials.map(|trials| {
|
||||||
let sampler = Arc::clone(sampler);
|
let sampler = Arc::clone(sampler);
|
||||||
let trials = Arc::clone(trials);
|
let trials = Arc::clone(trials);
|
||||||
|
let pruner = Arc::clone(pruner);
|
||||||
let factory: Arc<dyn Fn(u64) -> Trial + Send + Sync> = Arc::new(move |id| {
|
let factory: Arc<dyn Fn(u64) -> Trial + Send + Sync> = Arc::new(move |id| {
|
||||||
Trial::with_sampler(id, Arc::clone(&sampler), Arc::clone(&trials))
|
Trial::with_sampler(
|
||||||
|
id,
|
||||||
|
Arc::clone(&sampler),
|
||||||
|
Arc::clone(&trials),
|
||||||
|
Arc::clone(&pruner),
|
||||||
|
)
|
||||||
});
|
});
|
||||||
factory
|
factory
|
||||||
})
|
})
|
||||||
@@ -189,13 +198,14 @@ where
|
|||||||
V: 'static,
|
V: 'static,
|
||||||
{
|
{
|
||||||
let sampler: Arc<dyn Sampler> = Arc::new(sampler);
|
let sampler: Arc<dyn Sampler> = Arc::new(sampler);
|
||||||
|
let pruner: Arc<dyn Pruner> = Arc::new(pruner);
|
||||||
let completed_trials = Arc::new(RwLock::new(Vec::new()));
|
let completed_trials = Arc::new(RwLock::new(Vec::new()));
|
||||||
let trial_factory = Self::make_trial_factory(&sampler, &completed_trials);
|
let trial_factory = Self::make_trial_factory(&sampler, &completed_trials, &pruner);
|
||||||
|
|
||||||
Self {
|
Self {
|
||||||
direction,
|
direction,
|
||||||
sampler,
|
sampler,
|
||||||
pruner: Arc::new(pruner),
|
pruner,
|
||||||
completed_trials,
|
completed_trials,
|
||||||
next_trial_id: AtomicU64::new(0),
|
next_trial_id: AtomicU64::new(0),
|
||||||
trial_factory,
|
trial_factory,
|
||||||
@@ -207,7 +217,8 @@ where
|
|||||||
V: 'static,
|
V: 'static,
|
||||||
{
|
{
|
||||||
self.sampler = Arc::new(sampler);
|
self.sampler = Arc::new(sampler);
|
||||||
self.trial_factory = Self::make_trial_factory(&self.sampler, &self.completed_trials);
|
self.trial_factory =
|
||||||
|
Self::make_trial_factory(&self.sampler, &self.completed_trials, &self.pruner);
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Sets a new pruner for the study.
|
/// Sets a new pruner for the study.
|
||||||
@@ -215,8 +226,13 @@ where
|
|||||||
/// # Arguments
|
/// # Arguments
|
||||||
///
|
///
|
||||||
/// * `pruner` - The pruner to use for trial pruning.
|
/// * `pruner` - The pruner to use for trial pruning.
|
||||||
pub fn set_pruner(&mut self, pruner: impl Pruner + 'static) {
|
pub fn set_pruner(&mut self, pruner: impl Pruner + 'static)
|
||||||
|
where
|
||||||
|
V: 'static,
|
||||||
|
{
|
||||||
self.pruner = Arc::new(pruner);
|
self.pruner = Arc::new(pruner);
|
||||||
|
self.trial_factory =
|
||||||
|
Self::make_trial_factory(&self.sampler, &self.completed_trials, &self.pruner);
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Returns a reference to the study's pruner.
|
/// Returns a reference to the study's pruner.
|
||||||
@@ -288,12 +304,13 @@ where
|
|||||||
/// ```
|
/// ```
|
||||||
pub fn complete_trial(&self, mut trial: Trial, value: V) {
|
pub fn complete_trial(&self, mut trial: Trial, value: V) {
|
||||||
trial.set_complete();
|
trial.set_complete();
|
||||||
let completed = CompletedTrial::new(
|
let completed = CompletedTrial::with_intermediate_values(
|
||||||
trial.id(),
|
trial.id(),
|
||||||
trial.params().clone(),
|
trial.params().clone(),
|
||||||
trial.distributions().clone(),
|
trial.distributions().clone(),
|
||||||
trial.param_labels().clone(),
|
trial.param_labels().clone(),
|
||||||
value,
|
value,
|
||||||
|
trial.intermediate_values().to_vec(),
|
||||||
);
|
);
|
||||||
self.completed_trials.write().push(completed);
|
self.completed_trials.write().push(completed);
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -9,6 +9,7 @@ use crate::distribution::Distribution;
|
|||||||
use crate::error::{Error, Result};
|
use crate::error::{Error, Result};
|
||||||
use crate::param::ParamValue;
|
use crate::param::ParamValue;
|
||||||
use crate::parameter::{ParamId, Parameter};
|
use crate::parameter::{ParamId, Parameter};
|
||||||
|
use crate::pruner::Pruner;
|
||||||
use crate::sampler::{CompletedTrial, Sampler};
|
use crate::sampler::{CompletedTrial, Sampler};
|
||||||
use crate::types::TrialState;
|
use crate::types::TrialState;
|
||||||
|
|
||||||
@@ -36,6 +37,10 @@ pub struct Trial {
|
|||||||
sampler: Option<Arc<dyn Sampler>>,
|
sampler: Option<Arc<dyn Sampler>>,
|
||||||
/// Access to the history of completed trials (shared with Study).
|
/// Access to the history of completed trials (shared with Study).
|
||||||
history: Option<Arc<RwLock<Vec<CompletedTrial<f64>>>>>,
|
history: Option<Arc<RwLock<Vec<CompletedTrial<f64>>>>>,
|
||||||
|
/// Intermediate objective values reported at each step.
|
||||||
|
intermediate_values: Vec<(u64, f64)>,
|
||||||
|
/// The pruner used to decide whether to stop this trial early.
|
||||||
|
pruner: Option<Arc<dyn Pruner>>,
|
||||||
}
|
}
|
||||||
|
|
||||||
impl core::fmt::Debug for Trial {
|
impl core::fmt::Debug for Trial {
|
||||||
@@ -48,6 +53,8 @@ impl core::fmt::Debug for Trial {
|
|||||||
.field("param_labels", &self.param_labels)
|
.field("param_labels", &self.param_labels)
|
||||||
.field("has_sampler", &self.sampler.is_some())
|
.field("has_sampler", &self.sampler.is_some())
|
||||||
.field("has_history", &self.history.is_some())
|
.field("has_history", &self.history.is_some())
|
||||||
|
.field("intermediate_values", &self.intermediate_values)
|
||||||
|
.field("has_pruner", &self.pruner.is_some())
|
||||||
.finish()
|
.finish()
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
@@ -83,6 +90,8 @@ impl Trial {
|
|||||||
param_labels: HashMap::new(),
|
param_labels: HashMap::new(),
|
||||||
sampler: None,
|
sampler: None,
|
||||||
history: None,
|
history: None,
|
||||||
|
intermediate_values: Vec::new(),
|
||||||
|
pruner: None,
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -100,6 +109,7 @@ impl Trial {
|
|||||||
id: u64,
|
id: u64,
|
||||||
sampler: Arc<dyn Sampler>,
|
sampler: Arc<dyn Sampler>,
|
||||||
history: Arc<RwLock<Vec<CompletedTrial<f64>>>>,
|
history: Arc<RwLock<Vec<CompletedTrial<f64>>>>,
|
||||||
|
pruner: Arc<dyn Pruner>,
|
||||||
) -> Self {
|
) -> Self {
|
||||||
Self {
|
Self {
|
||||||
id,
|
id,
|
||||||
@@ -109,6 +119,8 @@ impl Trial {
|
|||||||
param_labels: HashMap::new(),
|
param_labels: HashMap::new(),
|
||||||
sampler: Some(sampler),
|
sampler: Some(sampler),
|
||||||
history: Some(history),
|
history: Some(history),
|
||||||
|
intermediate_values: Vec::new(),
|
||||||
|
pruner: Some(pruner),
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -159,6 +171,44 @@ impl Trial {
|
|||||||
&self.param_labels
|
&self.param_labels
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/// Reports an intermediate objective value at a given step.
|
||||||
|
///
|
||||||
|
/// Steps should be monotonically increasing (e.g., epoch number).
|
||||||
|
/// Duplicate steps overwrite the previous value.
|
||||||
|
pub fn report(&mut self, step: u64, value: f64) {
|
||||||
|
if let Some(entry) = self
|
||||||
|
.intermediate_values
|
||||||
|
.iter_mut()
|
||||||
|
.find(|(s, _)| *s == step)
|
||||||
|
{
|
||||||
|
entry.1 = value;
|
||||||
|
} else {
|
||||||
|
self.intermediate_values.push((step, value));
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Ask whether this trial should be pruned at the current step.
|
||||||
|
///
|
||||||
|
/// Returns `true` if the pruner recommends stopping this trial.
|
||||||
|
/// The caller should return `Err(TrialPruned)` from the objective.
|
||||||
|
#[must_use]
|
||||||
|
pub fn should_prune(&self) -> bool {
|
||||||
|
let (Some(pruner), Some(history)) = (&self.pruner, &self.history) else {
|
||||||
|
return false;
|
||||||
|
};
|
||||||
|
let Some(&(step, _)) = self.intermediate_values.last() else {
|
||||||
|
return false;
|
||||||
|
};
|
||||||
|
let history_guard = history.read();
|
||||||
|
pruner.should_prune(self.id, step, &self.intermediate_values, &history_guard)
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Returns all intermediate values reported so far.
|
||||||
|
#[must_use]
|
||||||
|
pub fn intermediate_values(&self) -> &[(u64, f64)] {
|
||||||
|
&self.intermediate_values
|
||||||
|
}
|
||||||
|
|
||||||
/// Sets the trial state to Complete.
|
/// Sets the trial state to Complete.
|
||||||
pub(crate) fn set_complete(&mut self) {
|
pub(crate) fn set_complete(&mut self) {
|
||||||
self.state = TrialState::Complete;
|
self.state = TrialState::Complete;
|
||||||
|
|||||||
Reference in New Issue
Block a user