From aba52f751d7c06d319c75486ed1579d4bed57aaf Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?R=C3=A9mi=20Lafage?= Date: Mon, 30 Sep 2024 21:44:18 +0200 Subject: [PATCH] Improve hot start API (#199) * Improve hot start API * Fix visibility * Use HotStartMode enum in EgorConfig --- ego/src/egor.rs | 26 +++++++++++++++++----- ego/src/lib.rs | 42 ++++++++++++++++++++++++++++++++--- ego/src/solver/egor_config.rs | 7 +++--- ego/src/utils/hot_start.rs | 41 +++++++++++++++++++++++++++++----- src/egor.rs | 4 ++-- 5 files changed, 101 insertions(+), 19 deletions(-) diff --git a/ego/src/egor.rs b/ego/src/egor.rs index 6624b8cc..dd1e572d 100644 --- a/ego/src/egor.rs +++ b/ego/src/egor.rs @@ -6,6 +6,7 @@ //! * Trust-region EGO optional activation //! * Infill criteria: EI, WB2, WB2S //! * Multi-point infill strategy (aka qEI) +//! * Warm/hot start //! //! See refences below. //! @@ -103,6 +104,7 @@ use crate::gpmix::mixint::*; use crate::types::*; use crate::EgorConfig; use crate::EgorState; +use crate::HotStartMode; use crate::{to_xtypes, EgorSolver}; use crate::{CheckpointingFrequency, HotStartCheckpoint}; @@ -213,12 +215,12 @@ impl Egor { let exec = Executor::new(self.fobj.clone(), self.solver.clone()); - let exec = if let Some(ext_iters) = self.solver.config.hot_start { + let exec = if self.solver.config.hot_start != HotStartMode::Disabled { let checkpoint = HotStartCheckpoint::new( ".checkpoints", "egor", CheckpointingFrequency::Always, - ext_iters, + self.solver.config.hot_start.clone(), ); exec.checkpointing(checkpoint) } else { @@ -423,7 +425,12 @@ mod tests { let _ = std::fs::remove_file(".checkpoints/egor.arg"); let n_iter = 1; let res = EgorBuilder::optimize(xsinx) - .configure(|config| config.max_iters(n_iter).seed(42).hot_start(Some(0))) + .configure(|config| { + config + .max_iters(n_iter) + .seed(42) + .hot_start(HotStartMode::Enabled) + }) .min_within(&array![[0.0, 25.0]]) .run() .expect("Egor should minimize"); @@ -432,7 +439,12 @@ mod tests { // without hostart we reach the same point let res = EgorBuilder::optimize(xsinx) - .configure(|config| config.max_iters(n_iter).seed(42).hot_start(None)) + .configure(|config| { + config + .max_iters(n_iter) + .seed(42) + .hot_start(HotStartMode::Disabled) + }) .min_within(&array![[0.0, 25.0]]) .run() .expect("Egor should minimize"); @@ -442,7 +454,11 @@ mod tests { // with hot start we continue let ext_iters = 3; let res = EgorBuilder::optimize(xsinx) - .configure(|config| config.seed(42).hot_start(Some(ext_iters))) + .configure(|config| { + config + .seed(42) + .hot_start(HotStartMode::ExtendedIters(ext_iters)) + }) .min_within(&array![[0.0, 25.0]]) .run() .expect("Egor should minimize"); diff --git a/ego/src/lib.rs b/ego/src/lib.rs index 33af59db..6e4ca1c5 100644 --- a/ego/src/lib.rs +++ b/ego/src/lib.rs @@ -11,7 +11,9 @@ //! * specify the initial doe, //! * parameterize internal optimization, //! * parameterize mixture of experts, -//! * save intermediate results and allow warm restart, +//! * save intermediate results and allow warm/hot restart, +//! * handling of mixed-integer variables +//! * activation of TREGO algorithm variation //! //! # Examples //! @@ -149,13 +151,47 @@ //! In the above example all GP with combinations of regression and correlation will be tested and the best combination for //! each modeled function will be retained. You can also simply specify `RegressionSpec::ALL` and `CorrelationSpec::ALL` to //! test all available combinations but remember that the more you test the slower it runs. +//! +//! * the TREGO algorithm described in \[[Diouane2023](#Diouane2023)\] can be activated +//! +//! ```no_run +//! # use egobox_ego::{EgorConfig, RegressionSpec, CorrelationSpec}; +//! # let egor_config = EgorConfig::default(); +//! egor_config.trego(true); +//! ``` +//! +//! * Intermediate results can be logged at each iteration when `outdir` directory is specified. +//! The following files : +//! * egor_config.json: Egor configuration, +//! * egor_initial_doe.npy: initial DOE (x, y) as numpy array, +//! * egor_doe.npy: DOE (x, y) as numpy array, +//! * egor_history.npy: best (x, y) wrt to iteration number as (n_iters, nx + ny) numpy array //! +//! ```no_run +//! # use egobox_ego::EgorConfig; +//! # let egor_config = EgorConfig::default(); +//! egor_config.outdir("./.output"); +//! ``` +//! If warm_start is set to `true`, the algorithm starts from the saved `egor_doe.npy` +//! +//! * Hot start checkpointing can be enabled with `hot_start` option specifying a number of +//! extra iterations beyond max iters. This mechanism allows to restart after an interruption +//! from the last saved checkpoint. While warm_start restart from saved doe for another max_iters +//! iterations, hot start allows to continue from the last saved optimizer state till max_iters +//! is reached with optinal extra iterations. +//! +//! ```no_run +//! # use egobox_ego::{EgorConfig, HotStartMode}; +//! # let egor_config = EgorConfig::default(); +//! egor_config.hot_start(HotStartMode::Enabled); +//! ``` +//! //! # Implementation notes //! //! * Mixture of experts and PLS dimension reduction is explained in \[[Bartoli2019](#Bartoli2019)\] //! * Parallel optimization is available through the selection of a qei strategy. See in \[[Ginsbourger2010](#Ginsbourger2010)\] //! * Mixed integer approach is implemented using continuous relaxation. See \[[Garrido2018](#Garrido2018)\] -//! * TREGO algorithm is enabled by default. See \[[Diouane2023](#Diouane2023)\] +//! * TREGO algorithm is not enabled by default. See \[[Diouane2023](#Diouane2023)\] //! //! # References //! @@ -210,7 +246,7 @@ pub use crate::gpmix::spec::{CorrelationSpec, RegressionSpec}; pub use crate::solver::*; pub use crate::types::*; pub use crate::utils::{ - find_best_result_index, Checkpoint, CheckpointingFrequency, HotStartCheckpoint, + find_best_result_index, Checkpoint, CheckpointingFrequency, HotStartCheckpoint, HotStartMode, }; mod optimizers; diff --git a/ego/src/solver/egor_config.rs b/ego/src/solver/egor_config.rs index 9685eb1b..1cd38f9e 100644 --- a/ego/src/solver/egor_config.rs +++ b/ego/src/solver/egor_config.rs @@ -1,6 +1,7 @@ //! Egor optimizer configuration. use crate::criteria::*; use crate::types::*; +use crate::HotStartMode; use egobox_moe::{CorrelationSpec, RegressionSpec}; use ndarray::Array1; use ndarray::Array2; @@ -81,7 +82,7 @@ pub struct EgorConfig { /// If true use `outdir` to retrieve and start from previous results pub(crate) warm_start: bool, /// If some enable checkpointing allowing to restart for given ext_iters number of iteration from last checkpointed iteration - pub(crate) hot_start: Option, + pub(crate) hot_start: HotStartMode, /// List of x types allowing the handling of discrete input variables pub(crate) xtypes: Vec, /// A random generator seed used to get reproductible results. @@ -111,7 +112,7 @@ impl Default for EgorConfig { target: f64::NEG_INFINITY, outdir: None, warm_start: false, - hot_start: None, + hot_start: HotStartMode::Disabled, xtypes: vec![], seed: None, trego: TregoConfig::default(), @@ -269,7 +270,7 @@ impl EgorConfig { } /// Whether checkpointing is enabled allowing hot start from previous checkpointed iteration if any - pub fn hot_start(mut self, hot_start: Option) -> Self { + pub fn hot_start(mut self, hot_start: HotStartMode) -> Self { self.hot_start = hot_start; self } diff --git a/ego/src/utils/hot_start.rs b/ego/src/utils/hot_start.rs index 1f6a4524..420fc497 100644 --- a/ego/src/utils/hot_start.rs +++ b/ego/src/utils/hot_start.rs @@ -1,17 +1,44 @@ pub use argmin::core::checkpointing::{Checkpoint, CheckpointingFrequency}; use argmin::core::Error; -use serde::{de::DeserializeOwned, Serialize}; +use serde::{de::DeserializeOwned, Deserialize, Serialize}; use std::fs::File; use std::io::{BufReader, BufWriter}; use std::path::PathBuf; use crate::EgorState; +/// An enum to specify hot start mode +#[derive(Clone, Eq, PartialEq, Debug, Hash, Default, Serialize, Deserialize)] +pub enum HotStartMode { + /// Hot start checkpoints are not saved + #[default] + Disabled, + /// Hot start checkpoints are saved and optionally used if it already exists + Enabled, + /// Hot start checkpoints are saved and optionally used if it already exists + /// and optimization is run with an extended iteration budget + ExtendedIters(u64), +} + +impl std::convert::From> for HotStartMode { + fn from(value: Option) -> Self { + if let Some(ext_iters) = value { + if ext_iters == 0 { + HotStartMode::Enabled + } else { + HotStartMode::ExtendedIters(ext_iters) + } + } else { + HotStartMode::Disabled + } + } +} + /// Handles saving a checkpoint to disk as a binary file. #[derive(Clone, Eq, PartialEq, Debug, Hash)] pub struct HotStartCheckpoint { /// Extended iteration number - pub extension_iters: u64, + pub mode: HotStartMode, /// Indicates how often a checkpoint is created pub frequency: CheckpointingFrequency, /// Directory where the checkpoints are saved to @@ -24,7 +51,7 @@ impl Default for HotStartCheckpoint { /// Create a default `HotStartCheckpoint` instance. fn default() -> HotStartCheckpoint { HotStartCheckpoint { - extension_iters: 0, + mode: HotStartMode::default(), frequency: CheckpointingFrequency::default(), directory: PathBuf::from(".checkpoints"), filename: PathBuf::from("egor.arg"), @@ -38,10 +65,10 @@ impl HotStartCheckpoint { directory: N, name: N, frequency: CheckpointingFrequency, - ext_iters: u64, + ext_iters: HotStartMode, ) -> Self { HotStartCheckpoint { - extension_iters: ext_iters, + mode: ext_iters, frequency, directory: PathBuf::from(directory.as_ref()), filename: PathBuf::from(format!("{}.arg", name.as_ref())), @@ -81,7 +108,9 @@ where let file = File::open(path)?; let reader = BufReader::new(file); let (solver, mut state): (_, EgorState<_>) = bincode::deserialize_from(reader)?; - state.extend_max_iters(self.extension_iters); + if let HotStartMode::ExtendedIters(n_iters) = self.mode { + state.extend_max_iters(n_iters); + } Ok(Some((solver, state))) } diff --git a/src/egor.rs b/src/egor.rs index a69d7e53..549b2b4a 100644 --- a/src/egor.rs +++ b/src/egor.rs @@ -138,7 +138,7 @@ pub(crate) fn to_specs(py: Python, xlimits: Vec>) -> PyResult /// warm_start (bool) /// Start by loading initial doe from directory /// -/// hot_start (int or None) +/// hot_start (int >= 0 or None) /// When hot_start>=0 saves optimizer state at each iteration and starts from a previous checkpoint /// if any for the given hot_start number of iterations beyond the max_iters nb of iterations. /// In an unstable environment were there can be crashes it allows to restart the optimization @@ -476,7 +476,7 @@ impl Egor { .n_optmod(self.n_optmod) .target(self.target) .warm_start(self.warm_start) - .hot_start(self.hot_start); + .hot_start(self.hot_start.into()); if let Some(doe) = doe { config = config.doe(doe); };