Skip to content

Commit

Permalink
Consolidate benchmark methods and give them a clearer naming system
Browse files Browse the repository at this point in the history 
Summary:
* No longer have "default" like "SOBOL+BOTORCH_MODULAR::default"; instead, construct method from components (e.g. "SOBOL+BOTORCH_MODULAR::SingleTaskGP_qLogNoisyExpectedImprovement") and *test* that the benchmarks match defaults. This has a few major benefits:

1) We won't test the same method on the same problem two or three times;

2) We will be able to do regression detection at the method level more easily, since there will be a 1:1 mapping between method name and functionality rather than many:many, and

3) tests will prevent benchmarks from drifting apart from default behavior (in MBM or CGS) over time.

* Remove non-Log EI benchmarks, since these have all been migrated.
* Remove ChooseGenerationStrategy benchmarks, replacing them with equivalent explicit MBM benchmarks where they don't already have them. Exceptions: Remove without replacing CGS benchmarks for MasterCook and IG Feature Selection, which do much better with Combinerator. `choose_gs_internal` chooses Combinerator, but `choose_generation_strategy` does not.  Similarly, remove `IGML_Surrogate:v1`, which does much better with a fully Bayesian model.
* Remove SOBOL+FULLYBAYESIAN, SOBOL+GPEI, and SOBOL+MAP_SAAS benchmarks, which were migrated to MBM earlier in this stack, and add benchmarks for the corresponding MBM methods where we don't already have them.
* Remove straight-through gradient benchmarks, since we don't use the method internally.
* Added a benchmark method registry like the benchmark problem registry to make it easier to track what's what
* TODO: revisit how these are parallelized

For a summary of changes, with notes and code pointers, see this spreadsheet:

https://docs.google.com/spreadsheets/d/1RQqkJKK-WWdZIluJNtwmeGwpbKt-J29r_pipk2T6e_s/edit#gid=1293825574

Differential Revision: https://internalfb.com/D49566866

fbshipit-source-id: b1faa552c91f782c504ee3998763cd168cb19e94
  • Loading branch information
@esantorella @facebook-github-bot
esantorella authored and facebook-github-bot committed Sep 27, 2023
1 parent 7a8069b commit c57f90c
Show file tree
Hide file tree
Showing 8 changed files with 199 additions and 346 deletions.
21 changes: 14 additions & 7 deletions ax/benchmark/benchmark.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -65,6 +65,19 @@ def compute_score_trace(
return score_trace.clip(min=0, max=100)


def _create_benchmark_experiment(
problem: BenchmarkProblemBase, method_name: str
) -> Experiment:
"""Creates an empty experiment for the given problem and method."""
return Experiment(
name=f"{problem.name}|{method_name}_{int(time())}",
search_space=problem.search_space,
optimization_config=problem.optimization_config,
tracking_metrics=problem.tracking_metrics,
runner=problem.runner,
)


def benchmark_replication(
problem: BenchmarkProblemBase,
method: BenchmarkMethod,
Expand All @@ -79,13 +92,7 @@ def benchmark_replication(
from `botorch.utils.sampling`.
"""

experiment = Experiment(
name=f"{problem.name}|{method.name}_{int(time())}",
search_space=problem.search_space,
optimization_config=problem.optimization_config,
tracking_metrics=problem.tracking_metrics,
runner=problem.runner,
)
experiment = _create_benchmark_experiment(problem=problem, method_name=method.name)

scheduler = Scheduler(
experiment=experiment,
Expand Down
30 changes: 1 addition & 29 deletions ax/benchmark/methods/choose_generation_strategy.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -3,32 +3,4 @@
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.

from typing import Optional

from ax.benchmark.benchmark_method import (
BenchmarkMethod,
get_sequential_optimization_scheduler_options,
)
from ax.benchmark.benchmark_problem import BenchmarkProblemBase
from ax.modelbridge.dispatch_utils import choose_generation_strategy
from ax.service.scheduler import SchedulerOptions


def get_choose_generation_strategy_method(
problem: BenchmarkProblemBase,
scheduler_options: Optional[SchedulerOptions] = None,
distribute_replications: bool = False,
) -> BenchmarkMethod:
generation_strategy = choose_generation_strategy(
search_space=problem.search_space,
optimization_config=problem.optimization_config,
num_trials=problem.num_trials,
)

return BenchmarkMethod(
name=f"ChooseGenerationStrategy::{problem.name}",
generation_strategy=generation_strategy,
scheduler_options=scheduler_options
or get_sequential_optimization_scheduler_options(),
distribute_replications=distribute_replications,
)
# TODO: remove this file in the next diff
56 changes: 1 addition & 55 deletions ax/benchmark/methods/gpei_and_moo.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -3,58 +3,4 @@
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.

from typing import Optional

from ax.benchmark.benchmark_method import (
BenchmarkMethod,
get_sequential_optimization_scheduler_options,
)
from ax.modelbridge.generation_strategy import GenerationStep, GenerationStrategy
from ax.modelbridge.registry import Models
from ax.service.scheduler import SchedulerOptions


def get_gpei_default(
scheduler_options: Optional[SchedulerOptions] = None,
) -> BenchmarkMethod:
generation_strategy = GenerationStrategy(
name="SOBOL+GPEI::default",
steps=[
GenerationStep(model=Models.SOBOL, num_trials=5, min_trials_observed=5),
GenerationStep(
model=Models.GPEI,
num_trials=-1,
max_parallelism=1,
),
],
)

return BenchmarkMethod(
name=generation_strategy.name,
generation_strategy=generation_strategy,
scheduler_options=scheduler_options
or get_sequential_optimization_scheduler_options(),
)


def get_moo_default(
scheduler_options: Optional[SchedulerOptions] = None,
) -> BenchmarkMethod:
generation_strategy = GenerationStrategy(
name="SOBOL+MOO::default",
steps=[
GenerationStep(model=Models.SOBOL, num_trials=5, min_trials_observed=5),
GenerationStep(
model=Models.MOO,
num_trials=-1,
max_parallelism=1,
),
],
)

return BenchmarkMethod(
name=generation_strategy.name,
generation_strategy=generation_strategy,
scheduler_options=scheduler_options
or get_sequential_optimization_scheduler_options(),
)
# TODO: delete this file in the next diff
223 changes: 14 additions & 209 deletions ax/benchmark/methods/modular_botorch.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,7 @@
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.

from typing import Any, Dict, Optional, Type
from typing import Dict, Optional, Type, Union

from ax.benchmark.benchmark_method import (
BenchmarkMethod,
Expand All @@ -13,228 +13,33 @@
from ax.modelbridge.registry import Models
from ax.models.torch.botorch_modular.surrogate import Surrogate
from ax.service.scheduler import SchedulerOptions
from ax.utils.common.constants import Keys
from botorch.acquisition.acquisition import AcquisitionFunction
from botorch.acquisition.monte_carlo import qNoisyExpectedImprovement
from botorch.acquisition.multi_objective.monte_carlo import (
qNoisyExpectedHypervolumeImprovement,
)
from botorch.models.fully_bayesian import SaasFullyBayesianSingleTaskGP
from botorch.models.gp_regression import FixedNoiseGP


def get_sobol_botorch_modular_fixed_noise_gp_qnei(
scheduler_options: Optional[SchedulerOptions] = None,
) -> BenchmarkMethod:
model_gen_kwargs = {
"model_gen_options": {
Keys.OPTIMIZER_KWARGS: {
"num_restarts": 50,
"raw_samples": 1024,
},
Keys.ACQF_KWARGS: {
"prune_baseline": True,
},
}
}

generation_strategy = GenerationStrategy(
name="SOBOL+BOTORCH_MODULAR::FixedNoiseGP_qNoisyExpectedImprovement",
steps=[
GenerationStep(model=Models.SOBOL, num_trials=5, min_trials_observed=5),
GenerationStep(
model=Models.BOTORCH_MODULAR,
num_trials=-1,
max_parallelism=1,
model_kwargs={
"surrogate": Surrogate(FixedNoiseGP),
"botorch_acqf_class": qNoisyExpectedImprovement,
},
model_gen_kwargs=model_gen_kwargs,
),
],
)

return BenchmarkMethod(
name=generation_strategy.name,
generation_strategy=generation_strategy,
scheduler_options=scheduler_options
or get_sequential_optimization_scheduler_options(),
)


def get_sobol_botorch_modular_fixed_noise_gp_qnehvi(
scheduler_options: Optional[SchedulerOptions] = None,
) -> BenchmarkMethod:
model_gen_kwargs = {
"model_gen_options": {
Keys.OPTIMIZER_KWARGS: {
"num_restarts": 50,
"raw_samples": 1024,
},
Keys.ACQF_KWARGS: {
"prune_baseline": True,
"qmc": True,
"mc_samples": 512,
},
}
}

generation_strategy = GenerationStrategy(
name="SOBOL+BOTORCH_MODULAR::FixedNoiseGP_qNoisyExpectedHypervolumeImprovement",
steps=[
GenerationStep(
model=Models.SOBOL,
num_trials=5,
min_trials_observed=5,
),
GenerationStep(
model=Models.BOTORCH_MODULAR,
num_trials=-1,
max_parallelism=1,
model_kwargs={
"surrogate": Surrogate(FixedNoiseGP),
"botorch_acqf_class": qNoisyExpectedHypervolumeImprovement,
},
model_gen_kwargs=model_gen_kwargs,
),
],
)

return BenchmarkMethod(
name=generation_strategy.name,
generation_strategy=generation_strategy,
scheduler_options=scheduler_options
or get_sequential_optimization_scheduler_options(),
)


def get_sobol_botorch_modular_saas_fully_bayesian_single_task_gp_qnei(
scheduler_options: Optional[SchedulerOptions] = None,
distribute_replications: bool = True,
) -> BenchmarkMethod: # noqa
return get_sobol_botorch_modular_saas_fully_bayesian_single_task_gp(
qNoisyExpectedImprovement,
scheduler_options=scheduler_options
or get_sequential_optimization_scheduler_options(),
distribute_replications=distribute_replications,
)


def get_sobol_botorch_modular_saas_fully_bayesian_single_task_gp(
botorch_acqf_class: Type[AcquisitionFunction],
scheduler_options: Optional[SchedulerOptions] = None,
distribute_replications: bool = True,
) -> BenchmarkMethod: # noqa
generation_strategy = GenerationStrategy(
name="SOBOL+BOTORCH_MODULAR::SaasFullyBayesianSingleTaskGP_"
+ botorch_acqf_class.__name__, # noqa
steps=[
GenerationStep(model=Models.SOBOL, num_trials=5, min_trials_observed=5),
GenerationStep(
model=Models.BOTORCH_MODULAR,
num_trials=-1,
max_parallelism=1,
model_kwargs={
"surrogate": Surrogate(
botorch_model_class=SaasFullyBayesianSingleTaskGP
),
"botorch_acqf_class": botorch_acqf_class,
},
),
],
)

return BenchmarkMethod(
name=generation_strategy.name,
generation_strategy=generation_strategy,
scheduler_options=scheduler_options
or get_sequential_optimization_scheduler_options(),
distribute_replications=distribute_replications,
)


def get_sobol_botorch_modular_saas_fully_bayesian_single_task_gp_qnehvi(
scheduler_options: Optional[SchedulerOptions] = None,
distribute_replications: bool = True,
) -> BenchmarkMethod: # noqa
generation_strategy = GenerationStrategy(
name="SOBOL+BOTORCH_MODULAR::SaasFullyBayesianSingleTaskGP_qNoisyExpectedHypervolumeImprovement", # noqa
steps=[
GenerationStep(
model=Models.SOBOL,
num_trials=5,
min_trials_observed=5,
),
GenerationStep(
model=Models.BOTORCH_MODULAR,
num_trials=-1,
max_parallelism=1,
model_kwargs={
"surrogate": Surrogate(
botorch_model_class=SaasFullyBayesianSingleTaskGP
),
"botorch_acqf_class": qNoisyExpectedHypervolumeImprovement,
},
),
],
)
return BenchmarkMethod(
name=generation_strategy.name,
generation_strategy=generation_strategy,
scheduler_options=scheduler_options
or get_sequential_optimization_scheduler_options(),
distribute_replications=distribute_replications,
)


def get_sobol_botorch_modular_default(
scheduler_options: Optional[SchedulerOptions] = None,
distribute_replications: bool = False,
) -> BenchmarkMethod:
generation_strategy = GenerationStrategy(
name="SOBOL+BOTORCH_MODULAR::default",
steps=[
GenerationStep(model=Models.SOBOL, num_trials=5, min_trials_observed=5),
GenerationStep(
model=Models.BOTORCH_MODULAR,
num_trials=-1,
max_parallelism=1,
),
],
)

return BenchmarkMethod(
name=generation_strategy.name,
generation_strategy=generation_strategy,
scheduler_options=scheduler_options
or get_sequential_optimization_scheduler_options(),
distribute_replications=distribute_replications,
)
from botorch.models.model import Model


def get_sobol_botorch_modular_acquisition(
model_cls: Type[Model],
acquisition_cls: Type[AcquisitionFunction],
acquisition_options: Optional[Dict[str, Any]] = None,
scheduler_options: Optional[SchedulerOptions] = None,
distribute_replications: bool = False,
name: Optional[str] = None,
) -> BenchmarkMethod:
model_kwargs: Dict[str, Union[Type[AcquisitionFunction], Surrogate]] = {
"botorch_acqf_class": acquisition_cls
}
model_kwargs["surrogate"] = Surrogate(model_cls)

generation_strategy = GenerationStrategy(
name=f"SOBOL+BOTORCH_MODULAR::{acquisition_cls.__name__}",
name=f"MBM::{model_cls.__name__}_{acquisition_cls.__name__}"
if name is None
else name,
steps=[
GenerationStep(
model=Models.SOBOL,
num_trials=5,
min_trials_observed=5,
),
GenerationStep(model=Models.SOBOL, num_trials=5, min_trials_observed=5),
GenerationStep(
model=Models.BOTORCH_MODULAR,
num_trials=-1,
max_parallelism=1,
model_kwargs={
"botorch_acqf_class": acquisition_cls,
"acquisition_options": acquisition_options,
},
model_kwargs=model_kwargs,
),
],
)
Expand Down
Loading

0 comments on commit c57f90c

Please sign in to comment.