Source code for ax.global_stopping.strategies.improvement

#!/usr/bin/env python3
# Copyright (c) Meta Platforms, Inc. and affiliates.
#
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.

from logging import Logger
from typing import Any, Dict, Optional, Tuple

import numpy as np
from ax.core.base_trial import TrialStatus
from ax.core.data import Data
from ax.core.experiment import Experiment
from ax.core.optimization_config import MultiObjectiveOptimizationConfig
from ax.core.trial import Trial
from ax.core.types import ComparisonOp
from ax.global_stopping.strategies.base import BaseGlobalStoppingStrategy
from ax.modelbridge.modelbridge_utils import observed_hypervolume
from ax.plot.pareto_utils import get_tensor_converter_model
from ax.utils.common.logger import get_logger
from ax.utils.common.typeutils import checked_cast

logger: Logger = get_logger(__name__)


[docs]class ImprovementGlobalStoppingStrategy(BaseGlobalStoppingStrategy): """ A stopping strategy which stops the optimization if there is no significant improvement over the iterations. For single-objective optimizations, this strategy stops the loop if the feasible (mean) objective has not improved over the past "window_size" iterations. In MOO loops, it stops the optimization loop if the hyper-volume of the pareto front has not improved in the past "window_size" iterations. """ def __init__( self, min_trials: int, window_size: int = 5, improvement_bar: float = 0.1 ) -> None: """ Initialize an improvement-based stopping strategy. Args: min_trials: Minimum number of trials before the stopping strategy kicks in. window_size: Number of recent trials to check the improvement in. improvement_bar: Threshold (in [0,1]) for considering relative improvement over the best point. """ super().__init__(min_trials=min_trials) self.window_size = window_size self.improvement_bar = improvement_bar self.hv_by_trial: Dict[int, float] = {}
[docs] def should_stop_optimization( self, experiment: Experiment, trial_to_check: Optional[int] = None, **kwargs: Dict[str, Any], ) -> Tuple[bool, str]: """ Check if the optimization has improved in the past "window_size" iterations. For single-objective optimization experiments, it will call _should_stop_single_objective() and for MOO experiments, it will call _should_stop_moo(). Before making either of these calls, this function carries out some sanity checks to handle obvious/invalid cases. Args: experiment: The experiment to apply the strategy on. trial_to_check: The trial in the experiment at which we want to check for stopping. If None, we check at the latest trial. Returns: A Tuple with a boolean determining whether the optimization should stop, and a str declaring the reason for stopping. """ if len(experiment.trials_by_status[TrialStatus.RUNNING]): message = "There are pending trials in the experiment." return False, message if len(experiment.trials_by_status[TrialStatus.COMPLETED]) == 0: message = "There are no completed trials yet." return False, message max_completed_trial = max( experiment.trial_indices_by_status[TrialStatus.COMPLETED] ) if trial_to_check is None: trial_to_check = max_completed_trial if trial_to_check > max_completed_trial: raise ValueError( "trial_to_check is larger than the total number of " f"trials (={max_completed_trial})." ) min_required_trials = max(self.min_trials, self.window_size) if trial_to_check < (min_required_trials - 1): stop = False message = ( "There are not enough completed trials to make a stop decision " f"(present: {trial_to_check+1}, required: {min_required_trials})." ) return stop, message if isinstance(experiment.optimization_config, MultiObjectiveOptimizationConfig): return self._should_stop_moo( experiment=experiment, trial_to_check=trial_to_check ) else: return self._should_stop_single_objective( experiment=experiment, trial_to_check=trial_to_check )
def _should_stop_moo( self, experiment: Experiment, trial_to_check: int ) -> Tuple[bool, str]: """ This is just the "should_stop_optimization" method of the class specialized to MOO experiments. It computes the (feasible) hypervolume of the pareto front at "trial_to_check" trial and "window_size" trials before, and suggest to stop the optimization if there is no significant improvement. Args: experiment: The experiment to apply the strategy on. trial_to_check: The trial in the experiment at which we want to check for stopping. If None, we check at the latest trial. Returns: A Tuple with a boolean determining whether the optimization should stop, and a str declaring the reason for stopping. """ reference_trial_index = trial_to_check - self.window_size + 1 data_df = experiment.fetch_data().df data_df_reference = data_df[data_df["trial_index"] <= reference_trial_index] data_df = data_df[data_df["trial_index"] <= trial_to_check] # Computing or retrieving HV at "window_size" iteration before if reference_trial_index in self.hv_by_trial: hv_reference = self.hv_by_trial[reference_trial_index] else: mb_reference = get_tensor_converter_model( experiment=experiment, data=Data(data_df_reference) ) hv_reference = observed_hypervolume(mb_reference) self.hv_by_trial[reference_trial_index] = hv_reference if hv_reference == 0: message = "The reference hypervolume is 0. Continue the optimization." return False, message # Computing HV at current trial mb = get_tensor_converter_model(experiment=experiment, data=Data(data_df)) hv = observed_hypervolume(mb) self.hv_by_trial[trial_to_check] = hv hv_improvement = (hv - hv_reference) / hv_reference stop = hv_improvement < self.improvement_bar if stop: message = ( f"The improvement in hypervolume in the past {self.window_size} " f"trials (={hv_improvement:.3f}) is less than {self.improvement_bar}." ) else: message = "" return stop, message def _should_stop_single_objective( self, experiment: Experiment, trial_to_check: int ) -> Tuple[bool, str]: """ This is the "should_stop_optimization" method of the class specialized to single-objective experiments. It computes the best feasible objective at "trial_to_check" trial and "window_size" trials before, and suggest to stop the trial if there is no significant improvement. Args: experiment: The experiment to apply the strategy on. trial_to_check: The trial in the experiment at which we want to check for stopping. If None, we check at the latest trial. Returns: A Tuple with a boolean determining whether the optimization should stop, and a str declaring the reason for stopping. """ objectives = [] is_feasible = [] for trial in experiment.trials_by_status[TrialStatus.COMPLETED]: if trial.index <= trial_to_check: tr = checked_cast(Trial, trial) objectives.append(tr.objective_mean) is_feasible.append(constraint_satisfaction(tr)) if experiment.optimization_config.objective.minimize: # pyre-ignore selector, mask_val = np.minimum, np.inf else: selector, mask_val = np.maximum, -np.inf # Replace objective value at infeasible iterations with mask_val masked_obj = np.where(is_feasible, objectives, mask_val) running_optimum = selector.accumulate(masked_obj) # Computing the interquartile for scaling the difference feasible_objectives = np.array(objectives)[is_feasible] if len(feasible_objectives) <= 1: message = "There are not enough feasible arms tried yet." return False, message q3, q1 = np.percentile(feasible_objectives, [75, 25]) iqr = q3 - q1 relative_improvement = np.abs( (running_optimum[-1] - running_optimum[-self.window_size]) / iqr ) stop = relative_improvement < self.improvement_bar if stop: message = ( f"The improvement in best objective in the past {self.window_size} " f"trials (={relative_improvement:.3f}) is less than " f"{self.improvement_bar}." ) else: message = "" return stop, message
[docs]def constraint_satisfaction(trial: Trial) -> bool: """ This function checks whether the outcome constraints of the optimization config of an experiment are satisfied in the given trial. Args: trial: A single-arm Trial at which we want to check the constraint. Returns: A boolean which is True iff all outcome constraints are satisifed. """ outcome_constraints = ( trial.experiment.optimization_config.outcome_constraints # pyre-ignore ) if len(outcome_constraints) == 0: return True df = trial.lookup_data().df for constraint in outcome_constraints: bound = constraint.bound metric_name = constraint.metric.name metric_data = df.loc[df["metric_name"] == metric_name] mean, sem = metric_data.iloc[0][["mean", "sem"]] if sem > 0.0: logger.warning( f"There is observation noise for metric {metric_name}. This may " "negatively affect the way we check constraint satisfaction." ) if constraint.op is ComparisonOp.LEQ: satisfied = mean <= bound else: satisfied = mean >= bound if not satisfied: return False return True