tutorials/modular_botax.ipynb (1,257 lines of code) (raw):
{
"cells": [
{
"cell_type": "code",
"execution_count": 1,
"metadata": {},
"outputs": [],
"source": [
"from typing import Any, Dict, Optional, Tuple, Type\n",
"\n",
"# Ax wrappers for BoTorch components\n",
"from ax.models.torch.botorch_modular.model import BoTorchModel\n",
"from ax.models.torch.botorch_modular.surrogate import Surrogate\n",
"from ax.models.torch.botorch_modular.list_surrogate import ListSurrogate\n",
"from ax.models.torch.botorch_modular.acquisition import Acquisition\n",
"\n",
"# Ax data tranformation layer\n",
"from ax.modelbridge.torch import TorchModelBridge\n",
"from ax.modelbridge.registry import Cont_X_trans, Y_trans, Models\n",
"\n",
"# Experiment examination utilities\n",
"from ax.service.utils.report_utils import exp_to_df\n",
"\n",
"# Test Ax objects\n",
"from ax.utils.testing.core_stubs import (\n",
" get_branin_experiment, \n",
" get_branin_data, \n",
" get_branin_experiment_with_multi_objective,\n",
" get_branin_data_multi_objective,\n",
")\n",
"\n",
"# BoTorch components\n",
"from botorch.models.model import Model\n",
"from botorch.models.gp_regression import FixedNoiseGP, SingleTaskGP\n",
"from botorch.acquisition.monte_carlo import qExpectedImprovement, qNoisyExpectedImprovement\n",
"from gpytorch.mlls.exact_marginal_log_likelihood import ExactMarginalLogLikelihood"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"# Setup and Usage of BoTorch Models in Ax\n",
"\n",
"Ax provides a set of flexible wrapper abstractions to mix-and-match BoTorch components like `Model` and `AcquisitionFunction` and combine them into a single `Model` object in Ax. The wrapper abstractions: `Surrogate`, `Acquisition`, and `BoTorchModel` – are located in `ax/models/torch/botorch_modular` directory and aim to encapsulate boilerplate code that interfaces between Ax and BoTorch. This functionality is in beta-release and still evolving.\n",
"\n",
"This tutorial walks through setting up a custom combination of BoTorch components in Ax in following steps:\n",
"\n",
"1. **Quick-start example of `BoTorchModel` use**\n",
"1. **`BoTorchModel` = `Surrogate` + `Acquisition` (overview)**\n",
" 1. Example with minimal options that uses the defaults\n",
" 2. Example showing all possible options\n",
" 3. Using a pre-constructed BoTorch Model (e.g. in research or development)\n",
" 4. Surrogate and Acquisition Q&A\n",
"2. **I know which Botorch Model and AcquisitionFunction I'd like to combine in Ax. How do set this up?**\n",
" 1. Making a `Surrogate` from BoTorch `Model`\n",
" 2. Using an arbitrary BoTorch `AcquisitionFunction` in Ax\n",
"3. **Using `Models.BOTORCH_MODULAR`** (convenience wrapper that enables storage and resumability)\n",
"4. **Utilizing `BoTorchModel` in generation strategies** (abstraction that allows to chain models together and use them in Ax Service API etc.)\n",
" 1. Specifying `pending_observations` to avoid the model re-suggesting points that are part of `RUNNING` or `ABANDONED` trials.\n",
"5. **Customizing a `Surrogate` or `Acquisition`** (for cases where existing subcomponent classes are not sufficient)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## 1. Quick-start example\n",
"\n",
"Here we set up a `BoTorchModel` with `SingleTaskGP` with `qNoisyExpectedImprovement`, one of the most popular combinations in Ax:"
]
},
{
"cell_type": "code",
"execution_count": 2,
"metadata": {},
"outputs": [
{
"name": "stderr",
"output_type": "stream",
"text": [
"[INFO 03-02 14:02:20] ax.core.experiment: The is_test flag has been set to True. This flag is meant purely for development and integration testing purposes. If you are running a live experiment, please set this flag to False\n"
]
}
],
"source": [
"experiment = get_branin_experiment(with_trial=True)\n",
"data = get_branin_data(trials=[experiment.trials[0]])"
]
},
{
"cell_type": "code",
"execution_count": 3,
"metadata": {},
"outputs": [
{
"name": "stderr",
"output_type": "stream",
"text": [
"[INFO 03-02 14:02:20] ax.modelbridge.transforms.standardize_y: Outcome branin is constant, within tolerance.\n"
]
}
],
"source": [
"# `Models` automatically selects a model + model bridge combination. \n",
"# For `BOTORCH_MODULAR`, it will select `BoTorchModel` and `TorchModelBridge`.\n",
"model_bridge_with_GPEI = Models.BOTORCH_MODULAR(\n",
" experiment=experiment,\n",
" data=data,\n",
" surrogate=Surrogate(SingleTaskGP), # Optional, will use default if unspecified\n",
" botorch_acqf_class=qNoisyExpectedImprovement, # Optional, will use default if unspecified\n",
")"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Now we can use this model to generate candidates (`gen`), predict outcome at a point (`predict`), or evaluate acquisition function value at a given point (`evaluate_acquisition_function`)."
]
},
{
"cell_type": "code",
"execution_count": 4,
"metadata": {},
"outputs": [
{
"data": {
"text/plain": [
"Arm(parameters={'x1': 10.0, 'x2': 15.0})"
]
},
"execution_count": 4,
"metadata": {
"bento_obj_id": "139741289584336"
},
"output_type": "execute_result"
}
],
"source": [
"generator_run = model_bridge_with_GPEI.gen(n=1)\n",
"generator_run.arms[0]"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"-----\n",
"Before you read the rest of this tutorial:\n",
"\n",
"- Note that the concept of ‘model’ is Ax is somewhat a misnomer; we use ['model'](https://ax.dev/docs/glossary.html#model) to refer to an optimization setup capable of producing candidate points for optimization (and often capable of being fit to data, with exception for quasi-random generators). See [Models documentation page](https://ax.dev/docs/models.html) for more information.\n",
"- Learn about `ModelBridge` in Ax, as users should rarely be interacting with a `Model` object directly (more about ModelBridge, a data transformation layer in Ax, [here](https://ax.dev/docs/models.html#deeper-dive-organization-of-the-modeling-stack))."
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## 2. BoTorchModel = Surrogate + Acquisition\n",
"\n",
"A `BoTorchModel` in Ax consists of two main subcomponents: a surrogate model and an acquisition function. A surrogate model is represented as an instance of Ax’s `Surrogate` class, which is a wrapper around BoTorch's `Model` class. The acquisition function is represented as an instance of Ax’s `Acquisition` class, a wrapper around BoTorch's `AcquisitionFunction` class."
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"### 2A. Example that uses defaults and requires no options\n",
"\n",
"BoTorchModel does not always require surrogate and acquisition specification. If instantiated without one or both components specified, defaults are selected based on properties of experiment and data (see Appendix 2 for auto-selection logic)."
]
},
{
"cell_type": "code",
"execution_count": 5,
"metadata": {},
"outputs": [],
"source": [
"# The surrogate is not specified, so it will be auto-selected\n",
"# during `model.fit`.\n",
"GPEI_model = BoTorchModel(botorch_acqf_class=qExpectedImprovement)\n",
"\n",
"# The acquisition class is not specified, so it will be \n",
"# auto-selected during `model.gen` or `model.evaluate_acquisition`\n",
"GPEI_model = BoTorchModel(surrogate=Surrogate(FixedNoiseGP))\n",
"\n",
"# Both the surrogate and acquisition class will be auto-selected.\n",
"GPEI_model = BoTorchModel()"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"### 2B. Example with all the options\n",
"Below are the full set of configurable settings of a `BoTorchModel` with their descriptions:"
]
},
{
"cell_type": "code",
"execution_count": 6,
"metadata": {},
"outputs": [],
"source": [
"model = BoTorchModel(\n",
" # Optional `Surrogate` specification to use instead of default\n",
" surrogate=Surrogate(\n",
" # BoTorch `Model` type\n",
" botorch_model_class=FixedNoiseGP,\n",
" # Optional, MLL class with which to optimize model parameters\n",
" mll_class=ExactMarginalLogLikelihood,\n",
" # Optional, dictionary of keyword arguments to underlying \n",
" # BoTorch `Model` constructor\n",
" model_options={}\n",
" ),\n",
" # Optional options to pass to auto-picked `Surrogate` if not\n",
" # specifying the `surrogate` argument\n",
" surrogate_options={},\n",
" \n",
" # Optional BoTorch `AcquisitionFunction` to use instead of default\n",
" botorch_acqf_class=qExpectedImprovement,\n",
" # Optional dict of keyword arguments, passed to the input \n",
" # constructor for the given BoTorch `AcquisitionFunction`\n",
" acquisition_options={},\n",
" # Optional Ax `Acquisition` subclass (if the given BoTorch\n",
" # `AcquisitionFunction` requires one, which is rare)\n",
" acquisition_class=None,\n",
" \n",
" # Less common model settings shown with default values, refer\n",
" # to `BoTorchModel` documentation for detail\n",
" refit_on_update=True,\n",
" refit_on_cv=False,\n",
" warm_start_refit=True,\n",
")"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## 2C. `Surrogate` from pre-instantiated BoTorch `Model`\n",
"\n",
"Alternatively, for BoTorch `Model`-s that require complex instantiation procedures (or is in development stage), leverage the `from_botorch` instantiation method of Surrogate:"
]
},
{
"cell_type": "code",
"execution_count": 7,
"metadata": {},
"outputs": [],
"source": [
"from_botorch_model = BoTorchModel(\n",
" surrogate=Surrogate.from_botorch(\n",
" # Pre-constructed BoTorch `Model` instance, with training data already set\n",
" model=..., \n",
" # Optional, MLL class with which to optimize model parameters\n",
" mll_class=ExactMarginalLogLikelihood,\n",
" )\n",
")"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## 2D. `Surrogate` and `Acquisition` Q&A\n",
"\n",
"**Why is the `surrogate` argument expected to be an instance, but `botorch_acqf_class` –– a class?** Because a BoTorch `AcquisitionFunction` object (and therefore its Ax wrapper, `Acquisition`) is ephemeral: it is constructed, immediately used, and destroyed during `BoTorchModel.gen`, so there is no reason to keep around an `Acquisition` instance. A `Surrogate`, on another hand, is kept in memory as long as its parent `BoTorchModel` is.\n",
"\n",
"**How to know when to use specify acquisition_class (and thereby a non-default Acquisition type) instead of just passing in botorch_acqf_class?** In short, custom `Acquisition` subclasses are needed when a given `AcquisitionFunction` in BoTorch needs some non-standard subcomponents or inputs (e.g. a custom BoTorch `AcquisitionObjective`). <TODO>\n",
"\n",
"**Please post any other questions you have to our dedicated issue on Github: https://github.com/facebook/Ax/issues/363.** This functionality is in beta-release and your feedback will be of great help to us!"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## 3. I know which Botorch `Model` and `AcquisitionFunction` I'd like to combine in Ax. How do set this up?"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"### 3a. Making a `Surrogate` from BoTorch `Model`:\n",
"Most models should work with base `Surrogate` in Ax, except for BoTorch `ModelListGP`, which works with `ListSurrogate`. `ModelListGP` is a special case because its purpose is to combine multiple sub-models into a single `Model` in BoTorch. It is most commonly used for multi-task optimization.\n",
"\n",
"If your `Model` is not a `ModelListGP`, the steps to set it up as a `Surrogate` are:\n",
"1. Implement a [`construct_inputs` class method](https://github.com/pytorch/botorch/blob/main/botorch/models/model.py#L143). The purpose of this method is to produce arguments to a particular model from a standardized set of inputs passed to BoTorch `Model`-s from [`Surrogate.construct`](https://github.com/facebook/Ax/blob/main/ax/models/torch/botorch_modular/surrogate.py#L148) in Ax. It should accept `TrainingData` and optionally other keyword arguments and produce a dictionary of arguments to `__init__` of the `Model`. See [`SingleTaskMultiFidelityGP.construct_inputs`](https://github.com/pytorch/botorch/blob/5b3172f3daa22f6ea2f6f4d1d0a378a9518dcd8d/botorch/models/gp_regression_fidelity.py#L131) for an example.\n",
"2. Pass any additional needed keyword arguments for the `Model` constructor (that cannot be constructed from `TrainingData` and other arguments to `construct_inputs`) via `model_options` argument to `Surrogate`."
]
},
{
"cell_type": "code",
"execution_count": 8,
"metadata": {},
"outputs": [],
"source": [
"from botorch.models.model import Model\n",
"from botorch.utils.containers import TrainingData\n",
"\n",
"class MyModelClass(Model):\n",
" \n",
" ... # Other contents of the `Model` type\n",
" \n",
" @classmethod\n",
" def construct_inputs(cls, training_data: TrainingData, **kwargs) -> Dict[str, Any]:\n",
" fidelity_features = kwargs.get(\"fidelity_features\")\n",
" if fidelity_features is None:\n",
" raise ValueError(f\"Fidelity features required for {cls.__name__}.\")\n",
"\n",
" return {\n",
" \"train_X\": training_data.X,\n",
" \"train_Y\": training_data.Y,\n",
" \"fidelity_features\": fidelity_features,\n",
" }\n",
"\n",
"surrogate = Surrogate(\n",
" botorch_model_class=MyModelClass, # Must implement `construct_inputs`\n",
" # Optional dict of additional keyword arguments to `MyModelClass`\n",
" model_options={},\n",
")"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"For a `ModelListGP`, the setup is similar, except that the surrogate is defined in terms of sub-models rather than one model. Both of the following options will work:"
]
},
{
"cell_type": "code",
"execution_count": 9,
"metadata": {},
"outputs": [],
"source": [
"class MyOtherModelClass(MyModelClass):\n",
" pass\n",
"\n",
"surrogate = ListSurrogate(\n",
" botorch_submodel_class_per_outcome={\n",
" \"metric_a\": MyModelClass, \n",
" \"metric_b\": MyOtherModelClass,\n",
" },\n",
" submodel_options_per_outcome={\"metric_a\": {}, \"metric_b\": {}},\n",
")"
]
},
{
"cell_type": "code",
"execution_count": 10,
"metadata": {},
"outputs": [],
"source": [
"surrogate = ListSurrogate(\n",
" # Shortcut if all submodels are the same type\n",
" botorch_submodel_class=MyModelClass,\n",
" # Shortcut if all submodel options are the same\n",
" submodel_options={},\n",
")"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"NOTE: if you run into a case where base `Surrogate` does not work with your BoTorch `Model`, please let us know in this Github issue: https://github.com/facebook/Ax/issues/363, so we can find the right solution and augment this tutorial."
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"### 3B. Using an arbitrary BoTorch `AcquisitionFunction` in Ax"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Steps to set up any `AcquisitionFunction` in Ax are:\n",
"1. Add an input constructor function to the correponding file in BoTorch: [botorch/acquisition/input_constructors.py](https://github.com/pytorch/botorch/blob/main/botorch/acquisition/input_constructors.py). The purpose of this method is to produce arguments to a acquisition function from a standardized set of inputs passed to BoTorch `AcquisitionFunction`-s from `Acquisition.__init__` in Ax. For example, see [`construct_inputs_qEHVI`](https://github.com/pytorch/botorch/blob/main/botorch/acquisition/input_constructors.py#L477), which creates a fairly complex set of arguments needed by `qExpectedHypervolumeImprovement` –– a popular multi-objective optimization acquisition function offered in Ax and BoTorch. \n",
" 1. Note that the new input construtor needs to be decorated with `@acqf_input_constructor(AcquisitionFunctionClass)`.\n",
"2. (Optional) If a given `AcquisitionFunction` requires specific options passed to the BoTorch `optimize_acqf`, it's possible to add default optimizer options for a given `AcquisitionFunction` to avoid always manually passing them via `acquisition_options`.\n",
"3. Specify the BoTorch `AcquisitionFunction` class as `botorch_acqf_class` to `BoTorchModel`\n",
"4. (Optional) Pass any additional keyword arguments to acquisition function constructor or to the optimizer function via `acquisition_options` argument to `BoTorchModel`."
]
},
{
"cell_type": "code",
"execution_count": 11,
"metadata": {},
"outputs": [
{
"data": {
"text/plain": [
"<ax.models.torch.botorch_modular.model.BoTorchModel at 0x7f180b1dc9d0>"
]
},
"execution_count": 11,
"metadata": {
"bento_obj_id": "139741242444240"
},
"output_type": "execute_result"
}
],
"source": [
"from ax.models.torch.botorch_modular.optimizer_argparse import optimizer_argparse\n",
"from botorch.acquisition.acquisition import AcquisitionFunction\n",
"from botorch.acquisition.input_constructors import acqf_input_constructor\n",
"from botorch.acquisition.objective import AcquisitionObjective\n",
"from torch import Tensor\n",
"\n",
"\n",
"class MyAcquisitionFunctionClass(AcquisitionFunction):\n",
" ... # Actual contents of the acquisition function class.\n",
"\n",
"\n",
"# 1. Add input constructor\n",
"@acqf_input_constructor(MyAcquisitionFunctionClass)\n",
"def construct_inputs_my_acqf(\n",
" model: Model,\n",
" training_data: TrainingData,\n",
" objective_thresholds: Tensor,\n",
" objective: Optional[AcquisitionObjective] = None,\n",
" **kwargs: Any,\n",
") -> Dict[str, Any]:\n",
" pass\n",
"\n",
"\n",
"# 2. Register default optimizer options\n",
"@optimizer_argparse.register(MyAcquisitionFunctionClass)\n",
"def _argparse_my_acqf(acqf: MyAcquisitionFunctionClass, sequential: bool = True) -> dict:\n",
" return {\"sequential\": sequential} # default to sequentially optimizing batches of queries\n",
"\n",
"\n",
"# 3-4. Specifying `botorch_acqf_class` and `acquisition_options`\n",
"BoTorchModel(\n",
" botorch_acqf_class=MyAcquisitionFunctionClass,\n",
" acquisition_options={\n",
" \"alpha\": 10 ** -6,\n",
" # The sub-dict by the key \"optimizer_options\" can be passed\n",
" # to propagate options to `optimize_acqf`, used in\n",
" # `Acquisition.optimize`, to add/override the default\n",
" # optimizer options registered above.\n",
" \"optimizer_options\": {\"sequential\": False},\n",
" },\n",
")"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"See section 2A for combining the resulting `Surrogate` instance and `Acquisition` type into a `BoTorchModel`. You can also leverage `Models.BOTORCH_MODULAR` for ease of use; more on it in section 4 below or in section 1 quick-start example."
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## 4. Using `Models.BOTORCH_MODULAR` and `Models.MOO_MODULAR`\n",
"\n",
"To simplify the instantiation of an Ax ModelBridge and its undelying Model, Ax provides a [`Models` registry enum](https://github.com/facebook/Ax/blob/main/ax/modelbridge/registry.py#L355). When calling entries of that enum (e.g. `Models.BOTORCH_MODULAR(experiment, data)`), the inputs are automatically distributed between a `Model` and a `ModelBridge` for a given setup. A call to a `Model` enum member yields a model bridge with an underlying model, ready for use to generate candidates.\n",
"\n",
"Here we use `Models.BOTORCH_MODULAR` to set up a model with all-default subcomponents:"
]
},
{
"cell_type": "code",
"execution_count": 12,
"metadata": {},
"outputs": [
{
"name": "stderr",
"output_type": "stream",
"text": [
"[INFO 03-02 14:02:21] ax.modelbridge.transforms.standardize_y: Outcome branin is constant, within tolerance.\n",
"/mnt/xarfuse/uid-450439/beabe8ad-seed-nspid4026531836_cgpid112231-ns-4026531840/gpytorch/utils/cholesky.py:40: NumericalWarning:\n",
"\n",
"A not p.d., added jitter of 1.0e-08 to the diagonal\n",
"\n"
]
},
{
"data": {
"text/plain": [
"GeneratorRun(1 arms, total weight 1.0)"
]
},
"execution_count": 12,
"metadata": {
"bento_obj_id": "139741263983472"
},
"output_type": "execute_result"
}
],
"source": [
"model_bridge_with_GPEI = Models.BOTORCH_MODULAR(\n",
" experiment=experiment, data=data,\n",
")\n",
"model_bridge_with_GPEI.gen(1)"
]
},
{
"cell_type": "code",
"execution_count": 13,
"metadata": {},
"outputs": [
{
"data": {
"text/plain": [
"botorch.acquisition.monte_carlo.qNoisyExpectedImprovement"
]
},
"execution_count": 13,
"metadata": {
"bento_obj_id": "139741607224336"
},
"output_type": "execute_result"
}
],
"source": [
"model_bridge_with_GPEI.model.botorch_acqf_class"
]
},
{
"cell_type": "code",
"execution_count": 14,
"metadata": {},
"outputs": [
{
"data": {
"text/plain": [
"botorch.models.gp_regression.FixedNoiseGP"
]
},
"execution_count": 14,
"metadata": {
"bento_obj_id": "139741605761040"
},
"output_type": "execute_result"
}
],
"source": [
"model_bridge_with_GPEI.model.surrogate.botorch_model_class"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"And here we use `Models.MOO_MODULAR` (analogue of `Models.BOTORCH_MODULAR`, except set up with multi-objective model bridge) to set up a model for multi-objective optimization:"
]
},
{
"cell_type": "code",
"execution_count": 15,
"metadata": {},
"outputs": [
{
"name": "stderr",
"output_type": "stream",
"text": [
"[INFO 03-02 14:02:22] ax.core.experiment: The is_test flag has been set to True. This flag is meant purely for development and integration testing purposes. If you are running a live experiment, please set this flag to False\n",
"[INFO 03-02 14:02:22] ax.modelbridge.transforms.standardize_y: Outcome branin_a is constant, within tolerance.\n",
"[INFO 03-02 14:02:22] ax.modelbridge.transforms.standardize_y: Outcome branin_b is constant, within tolerance.\n"
]
},
{
"data": {
"text/plain": [
"GeneratorRun(1 arms, total weight 1.0)"
]
},
"execution_count": 15,
"metadata": {
"bento_obj_id": "139741370925744"
},
"output_type": "execute_result"
}
],
"source": [
"model_bridge_with_EHVI = Models.MOO_MODULAR(\n",
" experiment=get_branin_experiment_with_multi_objective(has_objective_thresholds=True, with_batch=True),\n",
" data=get_branin_data_multi_objective(),\n",
")\n",
"model_bridge_with_EHVI.gen(1)"
]
},
{
"cell_type": "code",
"execution_count": 16,
"metadata": {},
"outputs": [
{
"data": {
"text/plain": [
"botorch.acquisition.multi_objective.monte_carlo.qNoisyExpectedHypervolumeImprovement"
]
},
"execution_count": 16,
"metadata": {
"bento_obj_id": "139741606167568"
},
"output_type": "execute_result"
}
],
"source": [
"model_bridge_with_EHVI.model.botorch_acqf_class"
]
},
{
"cell_type": "code",
"execution_count": 17,
"metadata": {},
"outputs": [
{
"data": {
"text/plain": [
"botorch.models.gp_regression.FixedNoiseGP"
]
},
"execution_count": 17,
"metadata": {
"bento_obj_id": "139741605761040"
},
"output_type": "execute_result"
}
],
"source": [
"model_bridge_with_EHVI.model.surrogate.botorch_model_class"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Furthermore, the quick-start example at the top of this tutorial shows how to specify surrogate and acquisition subcomponents to `Models.BOTORCH_MODULAR`. "
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## 5. Utilizing `BoTorchModel` in generation strategies\n",
"\n",
"Generation strategy is a key concept in Ax, enabling use of Service API (a.k.a. `AxClient`) and many other higher-level abstractions. A `GenerationStrategy` allows to chain multiple models in Ax and thereby automate candidate generation. Refer to the \"Generation Strategy\" tutorial for more detail in generation strategies.\n",
"\n",
"An example generation stategy with the modular `BoTorchModel` would look like this:"
]
},
{
"cell_type": "code",
"execution_count": 18,
"metadata": {},
"outputs": [],
"source": [
"from ax.modelbridge.generation_strategy import GenerationStep, GenerationStrategy\n",
"from botorch.acquisition import UpperConfidenceBound\n",
"from ax.modelbridge.modelbridge_utils import get_pending_observation_features\n",
"\n",
"gs = GenerationStrategy(\n",
" steps=[\n",
" GenerationStep( # Initialization step\n",
" # Which model to use for this step\n",
" model=Models.SOBOL,\n",
" # How many generator runs (each of which is then made a trial) \n",
" # to produce with this step\n",
" num_trials=5,\n",
" # How many trials generated from this step must be `COMPLETED` \n",
" # before the next one\n",
" min_trials_observed=5, \n",
" ),\n",
" GenerationStep( # BayesOpt step\n",
" model=Models.BOTORCH_MODULAR,\n",
" # No limit on how many generator runs will be produced\n",
" num_trials=-1,\n",
" model_kwargs={ # Kwargs to pass to `BoTorchModel.__init__`\n",
" \"surrogate\": Surrogate(SingleTaskGP),\n",
" \"botorch_acqf_class\": qNoisyExpectedImprovement,\n",
" },\n",
" )\n",
" ]\n",
")"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Set up an experiment and generate 10 trials in it, adding synthetic data to experiment after each one:"
]
},
{
"cell_type": "code",
"execution_count": 19,
"metadata": {},
"outputs": [
{
"name": "stderr",
"output_type": "stream",
"text": [
"[INFO 03-02 14:02:22] ax.core.experiment: The is_test flag has been set to True. This flag is meant purely for development and integration testing purposes. If you are running a live experiment, please set this flag to False\n"
]
},
{
"data": {
"text/plain": [
"SearchSpace(parameters=[RangeParameter(name='x1', parameter_type=FLOAT, range=[-5.0, 10.0]), RangeParameter(name='x2', parameter_type=FLOAT, range=[0.0, 15.0])], parameter_constraints=[])"
]
},
"execution_count": 19,
"metadata": {
"bento_obj_id": "139741242475472"
},
"output_type": "execute_result"
}
],
"source": [
"experiment = get_branin_experiment(minimize=True)\n",
"\n",
"assert len(experiment.trials) == 0\n",
"experiment.search_space"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## 5a. Specifying `pending_observations`\n",
"Note that it's important to **specify pending observations** to the call to `gen` to avoid getting the same points re-suggested. Without `pending_observations` argument, Ax models are not aware of points that should be excluded from generation. Points are considered \"pending\" when they belong to `STAGED`, `RUNNING`, or `ABANDONED` trials (with the latter included so model does not re-suggest points that are considered \"bad\" and should not be re-suggested).\n",
"\n",
"If the call to `get_pending_obervation_features` becomes slow in your setup (since it performs data-fetching etc.), you can opt for `get_pending_observation_features_based_on_trial_status` (also from `ax.modelbridge.modelbridge_utils`), but note the limitations of that utility (detailed in its docstring)."
]
},
{
"cell_type": "code",
"execution_count": 20,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"Completed trial #0, suggested by Sobol.\n",
"Completed trial #1, suggested by Sobol.\n",
"Completed trial #2, suggested by Sobol.\n",
"Completed trial #3, suggested by Sobol.\n",
"Completed trial #4, suggested by Sobol.\n",
"Completed trial #5, suggested by BoTorch.\n",
"Completed trial #6, suggested by BoTorch.\n",
"Completed trial #7, suggested by BoTorch.\n",
"Completed trial #8, suggested by BoTorch.\n",
"Completed trial #9, suggested by BoTorch.\n"
]
}
],
"source": [
"for _ in range(10):\n",
" # Produce a new generator run and attach it to experiment as a trial\n",
" generator_run = gs.gen(\n",
" experiment=experiment, \n",
" n=1, \n",
" pending_observations=get_pending_observation_features(experiment=experiment),\n",
" )\n",
" trial = experiment.new_trial(generator_run)\n",
" \n",
" # Mark the trial as 'RUNNING' so we can mark it 'COMPLETED' later\n",
" trial.mark_running(no_runner_required=True)\n",
" \n",
" # Attach data for the new trial and mark it 'COMPLETED'\n",
" experiment.attach_data(get_branin_data(trials=[trial]))\n",
" trial.mark_completed()\n",
" \n",
" print(f\"Completed trial #{trial.index}, suggested by {generator_run._model_key}.\")"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Now we examine the experiment and observe the trials that were added to it and produced by the generation strategy:"
]
},
{
"cell_type": "code",
"execution_count": 21,
"metadata": {},
"outputs": [
{
"data": {
"application/vnd.dataresource+json": {
"data": [
{
"arm_name": "0_0",
"branin": 207.5594261918,
"generation_method": "Sobol",
"index": 0,
"trial_index": 0,
"trial_status": "COMPLETED",
"x1": -4.0517528355,
"x2": 0.3050024249
},
{
"arm_name": "1_0",
"branin": 15.2331241597,
"generation_method": "Sobol",
"index": 1,
"trial_index": 1,
"trial_status": "COMPLETED",
"x1": 2.592086019,
"x2": 6.4056294272
},
{
"arm_name": "2_0",
"branin": 6.70614789,
"generation_method": "Sobol",
"index": 2,
"trial_index": 2,
"trial_status": "COMPLETED",
"x1": 4.1784402775,
"x2": 2.8651705664
},
{
"arm_name": "3_0",
"branin": 65.5882072176,
"generation_method": "Sobol",
"index": 3,
"trial_index": 3,
"trial_status": "COMPLETED",
"x1": 8.8759032032,
"x2": 10.0373605639
},
{
"arm_name": "4_0",
"branin": 13.546869493,
"generation_method": "Sobol",
"index": 4,
"trial_index": 4,
"trial_status": "COMPLETED",
"x1": 1.2252612365,
"x2": 4.7846784303
},
{
"arm_name": "5_0",
"branin": 10.9599150042,
"generation_method": "BoTorch",
"index": 5,
"trial_index": 5,
"trial_status": "COMPLETED",
"x1": 10,
"x2": 0
},
{
"arm_name": "6_0",
"branin": 8.0161501028,
"generation_method": "BoTorch",
"index": 6,
"trial_index": 6,
"trial_status": "COMPLETED",
"x1": -3.1273412034,
"x2": 15
},
{
"arm_name": "7_0",
"branin": 155.9206002197,
"generation_method": "BoTorch",
"index": 7,
"trial_index": 7,
"trial_status": "COMPLETED",
"x1": 2.8069694499,
"x2": 15
},
{
"arm_name": "8_0",
"branin": 53.0342883398,
"generation_method": "BoTorch",
"index": 8,
"trial_index": 8,
"trial_status": "COMPLETED",
"x1": -5,
"x2": 10.8381700636
},
{
"arm_name": "9_0",
"branin": 14.2394033154,
"generation_method": "BoTorch",
"index": 9,
"trial_index": 9,
"trial_status": "COMPLETED",
"x1": 8.055378798,
"x2": 4.0433708202
}
],
"schema": {
"fields": [
{
"name": "index",
"type": "integer"
},
{
"name": "branin",
"type": "number"
},
{
"name": "trial_index",
"type": "integer"
},
{
"name": "arm_name",
"type": "string"
},
{
"name": "x1",
"type": "number"
},
{
"name": "x2",
"type": "number"
},
{
"name": "trial_status",
"type": "string"
},
{
"name": "generation_method",
"type": "string"
}
],
"pandas_version": "0.20.0",
"primaryKey": [
"index"
]
}
},
"text/html": [
"<div>\n",
"<style scoped>\n",
" .dataframe tbody tr th:only-of-type {\n",
" vertical-align: middle;\n",
" }\n",
"\n",
" .dataframe tbody tr th {\n",
" vertical-align: top;\n",
" }\n",
"\n",
" .dataframe thead th {\n",
" text-align: right;\n",
" }\n",
"</style>\n",
"<table border=\"1\" class=\"dataframe\">\n",
" <thead>\n",
" <tr style=\"text-align: right;\">\n",
" <th></th>\n",
" <th>branin</th>\n",
" <th>trial_index</th>\n",
" <th>arm_name</th>\n",
" <th>x1</th>\n",
" <th>x2</th>\n",
" <th>trial_status</th>\n",
" <th>generation_method</th>\n",
" </tr>\n",
" </thead>\n",
" <tbody>\n",
" <tr>\n",
" <th>0</th>\n",
" <td>207.559426</td>\n",
" <td>0</td>\n",
" <td>0_0</td>\n",
" <td>-4.051753</td>\n",
" <td>0.305002</td>\n",
" <td>COMPLETED</td>\n",
" <td>Sobol</td>\n",
" </tr>\n",
" <tr>\n",
" <th>1</th>\n",
" <td>15.233124</td>\n",
" <td>1</td>\n",
" <td>1_0</td>\n",
" <td>2.592086</td>\n",
" <td>6.405629</td>\n",
" <td>COMPLETED</td>\n",
" <td>Sobol</td>\n",
" </tr>\n",
" <tr>\n",
" <th>2</th>\n",
" <td>6.706148</td>\n",
" <td>2</td>\n",
" <td>2_0</td>\n",
" <td>4.178440</td>\n",
" <td>2.865171</td>\n",
" <td>COMPLETED</td>\n",
" <td>Sobol</td>\n",
" </tr>\n",
" <tr>\n",
" <th>3</th>\n",
" <td>65.588207</td>\n",
" <td>3</td>\n",
" <td>3_0</td>\n",
" <td>8.875903</td>\n",
" <td>10.037361</td>\n",
" <td>COMPLETED</td>\n",
" <td>Sobol</td>\n",
" </tr>\n",
" <tr>\n",
" <th>4</th>\n",
" <td>13.546869</td>\n",
" <td>4</td>\n",
" <td>4_0</td>\n",
" <td>1.225261</td>\n",
" <td>4.784678</td>\n",
" <td>COMPLETED</td>\n",
" <td>Sobol</td>\n",
" </tr>\n",
" <tr>\n",
" <th>5</th>\n",
" <td>10.959915</td>\n",
" <td>5</td>\n",
" <td>5_0</td>\n",
" <td>10.000000</td>\n",
" <td>0.000000</td>\n",
" <td>COMPLETED</td>\n",
" <td>BoTorch</td>\n",
" </tr>\n",
" <tr>\n",
" <th>6</th>\n",
" <td>8.016150</td>\n",
" <td>6</td>\n",
" <td>6_0</td>\n",
" <td>-3.127341</td>\n",
" <td>15.000000</td>\n",
" <td>COMPLETED</td>\n",
" <td>BoTorch</td>\n",
" </tr>\n",
" <tr>\n",
" <th>7</th>\n",
" <td>155.920600</td>\n",
" <td>7</td>\n",
" <td>7_0</td>\n",
" <td>2.806969</td>\n",
" <td>15.000000</td>\n",
" <td>COMPLETED</td>\n",
" <td>BoTorch</td>\n",
" </tr>\n",
" <tr>\n",
" <th>8</th>\n",
" <td>53.034288</td>\n",
" <td>8</td>\n",
" <td>8_0</td>\n",
" <td>-5.000000</td>\n",
" <td>10.838170</td>\n",
" <td>COMPLETED</td>\n",
" <td>BoTorch</td>\n",
" </tr>\n",
" <tr>\n",
" <th>9</th>\n",
" <td>14.239403</td>\n",
" <td>9</td>\n",
" <td>9_0</td>\n",
" <td>8.055379</td>\n",
" <td>4.043371</td>\n",
" <td>COMPLETED</td>\n",
" <td>BoTorch</td>\n",
" </tr>\n",
" </tbody>\n",
"</table>\n",
"</div>"
],
"text/plain": [
" branin trial_index arm_name ... x2 trial_status generation_method\n",
"0 207.559426 0 0_0 ... 0.305002 COMPLETED Sobol\n",
"1 15.233124 1 1_0 ... 6.405629 COMPLETED Sobol\n",
"2 6.706148 2 2_0 ... 2.865171 COMPLETED Sobol\n",
"3 65.588207 3 3_0 ... 10.037361 COMPLETED Sobol\n",
"4 13.546869 4 4_0 ... 4.784678 COMPLETED Sobol\n",
"5 10.959915 5 5_0 ... 0.000000 COMPLETED BoTorch\n",
"6 8.016150 6 6_0 ... 15.000000 COMPLETED BoTorch\n",
"7 155.920600 7 7_0 ... 15.000000 COMPLETED BoTorch\n",
"8 53.034288 8 8_0 ... 10.838170 COMPLETED BoTorch\n",
"9 14.239403 9 9_0 ... 4.043371 COMPLETED BoTorch\n",
"\n",
"[10 rows x 7 columns]"
]
},
"execution_count": 21,
"metadata": {
"bento_obj_id": "139741241071312"
},
"output_type": "execute_result"
}
],
"source": [
"exp_to_df(experiment)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## 6. Customizing a `Surrogate` or `Acquisition`\n",
"\n",
"We expect the base `Surrogate`, `ListSurrogate`, and `Acquisition` classes to work with most BoTorch components, but there could be a case where you would need to subclass one of aforementioned abstractions to handle a given BoTorch component. If you run into a case like this, feel free to open an issue on our [Github issues page](https://github.com/facebook/Ax/issues) –– it would be very useful for us to know \n",
"\n",
"One such example would be a need for a custom `AcquisitionObjective` or for a custom acquisition function optimization utility. To subclass `Acquisition` accordingly, one would override the `get_botorch_objective` method:"
]
},
{
"cell_type": "code",
"execution_count": 22,
"metadata": {},
"outputs": [],
"source": [
"class CustomObjectiveAcquisition(Acquisition):\n",
" \n",
" def get_botorch_objective(\n",
" self,\n",
" botorch_acqf_class: Type[AcquisitionFunction],\n",
" model: Model,\n",
" objective_weights: Tensor,\n",
" objective_thresholds: Optional[Tensor] = None,\n",
" outcome_constraints: Optional[Tuple[Tensor, Tensor]] = None,\n",
" X_observed: Optional[Tensor] = None,\n",
" ) -> AcquisitionObjective:\n",
" ... # Produce the desired `AcquisitionObjective` instead of the default"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Then to use the new subclass in `BoTorchModel`, just specify `acquisition_class` argument along with `botorch_acqf_class` (to `BoTorchModel` directly or to `Models.BOTORCH_MODULAR`, which just passes the relevant arguments to `BoTorchModel` under the hood, as discussed in section 4):"
]
},
{
"cell_type": "code",
"execution_count": 23,
"metadata": {},
"outputs": [
{
"name": "stderr",
"output_type": "stream",
"text": [
"[INFO 03-02 14:02:24] ax.modelbridge.transforms.standardize_y: Outcome branin is constant, within tolerance.\n"
]
},
{
"data": {
"text/plain": [
"<ax.modelbridge.torch.TorchModelBridge at 0x7f180b0fa0d0>"
]
},
"execution_count": 23,
"metadata": {
"bento_obj_id": "139741241516240"
},
"output_type": "execute_result"
}
],
"source": [
"Models.BOTORCH_MODULAR(\n",
" experiment=experiment, \n",
" data=data,\n",
" acquisition_class=CustomObjectiveAcquisition,\n",
" botorch_acqf_class=MyAcquisitionFunctionClass,\n",
")"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"To use a custom `Surrogate` subclass, pass the `surrogate` argument of that type:\n",
"```\n",
"Models.BOTORCH_MODULAR(\n",
" experiment=experiment, \n",
" data=data,\n",
" surrogate=CustomSurrogate(botorch_model_class=MyModelClass),\n",
")\n",
"```"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"------"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Appendix 1: Methods available on `BoTorchModel`\n",
"\n",
"Note that usually all these methods are used through `ModelBridge` –– a convertion and transformation layer that adapts Ax abstractions to inputs required by the given model.\n",
"\n",
"**Core methods on `BoTorchModel`:**\n",
"* `fit` selects a surrogate if needed and fits the surrogate model to data via `Surrogate.fit`,\n",
"* `predict` estimates metric values at a given point via `Surrogate.predict`,\n",
"* `gen` instantiates an acquisition function via `Acquisition.__init__` and optimizes it to generate candidates.\n",
"\n",
"**Other methods on `BoTorchModel`:**\n",
"* `update` updates surrogate model with training data and optionally reoptimizes model parameters via `Surrogate.update`,\n",
"* `cross_validate` re-fits the surrogate model to subset of training data and makes predictions for test data,\n",
"* `evaluate_acquisition_function` instantiates an acquisition function and evaluates it for a given point.\n",
"------\n"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Appendix 2: Default surrogate models and acquisition functions\n",
"\n",
"By default, the chosen surrogate model will be:\n",
"* if fidelity parameters are present in search space: `FixedNoiseMultiFidelityGP` (if [SEM](https://ax.dev/docs/glossary.html#sem)s are known on observations) and `SingleTaskMultiFidelityGP` (if variance unknown and needs to be inferred),\n",
"* if task parameters are present: a set of `FixedNoiseMultiTaskGP` (if known variance) or `MultiTaskGP` (if unknown variance), wrapped in a `ModelListGP` and each modeling one task,\n",
"* `FixedNoiseGP` (known variance) and `SingleTaskGP` (unknown variance) otherwise.\n",
"\n",
"The chosen acquisition function will be:\n",
"* for multi-objective settings: `qExpectedHypervolumeImprovement`,\n",
"* `qExpectedImprovement` (known variance) and `qNoisyExpectedImprovement` (unknown variance) otherwise.\n",
"----"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Appendix 3: Handling storage errors that arise from objects that don't have serialization logic in A\n",
"\n",
"Attempting to store a generator run produced via `Models.BOTORCH_MODULAR` instance that included options without serization logic with will produce an error like: `\"Object <SomeAcquisitionOption object> passed to 'object_to_json' (of type <class SomeAcquisitionOption'>) is not registered with a corresponding encoder in ENCODER_REGISTRY.\"`"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"The two options for handling this error are:\n",
"1. disabling storage of `BoTorchModel`'s options by passing `no_model_options_storage=True` to `Models.BOTORCH_MODULAR(...)` call –– this will prevent model options from being stored on the generator run, so a generator run can be saved but cannot be used to restore the model that produced it,\n",
"2. specifying serialization logic for a given object that needs to occur among the `Model` or `AcquisitionFunction` options. Tutorial for this is in the works, but in the meantime you can [post an issue on the Ax GitHub](https://github.com/facebook/Ax/issues) to get help with this."
]
}
],
"metadata": {
"kernelspec": {
"name": "python3",
"display_name": "python3"
}
},
"nbformat": 4,
"nbformat_minor": 5
}