tutorials/botorch_modular.ipynb (627 lines of code) (raw):
{
"cells": [
{
"cell_type": "markdown",
"metadata": {
"code_folding": [],
"customInput": null,
"hidden_ranges": [],
"originalKey": "59a44903-3946-4c41-abc4-1c915eefeac4",
"showInput": false
},
"source": [
"# Tutorial notebook for modular `BoTorchModel` customization\n",
"\n",
"NOTE: The functionality in this tutorial is still in its alpha stages.\n",
"\n",
"Contents:\n",
"1. Overview of modular `BoTorchModel`\n",
"2. `BoTorchModel` instantiation\n",
"3. Use a custom BoTorch `AcquisitionFunction`\n",
" 1. [Path 1] Use the default Ax `Acquisition` class\n",
" 2. [Path 2] Create a custom Ax `Acquisition` subclass\n",
" 3. Set up storage for the new setup\n",
"4. Details of `BoTorchModel` Subcomponent Classes"
]
},
{
"cell_type": "markdown",
"metadata": {
"code_folding": [],
"collapsed": true,
"customInput": null,
"hidden_ranges": [],
"originalKey": "701d6c9c-506c-43b8-b065-18b2a1f39271",
"showInput": false
},
"source": [
"# Overview of modular `BoTorchModel`\n",
"\n",
"**`BoTorchModel` = `Surrogate` + `Acquisition`**\n",
"\n",
"A `BoTorchModel` 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](https://github.com/facebook/Ax/blob/main/ax/models/torch/botorch_modular/surrogate.py), which is a wrapper around [BoTorch's `Model` class](https://github.com/pytorch/botorch/blob/main/botorch/models/model.py). The acquisition function is represented as an instance of [Ax’s `Acquisition` class](https://github.com/facebook/Ax/blob/main/ax/models/torch/botorch_modular/acquisition.py), a wrapper around [BoTorch's `AcquisitionFunction` class](https://github.com/pytorch/botorch/blob/main/botorch/acquisition/acquisition.py). These two subcomponents are described in greater detail at the bottom of this tutorial.\n",
"\n",
"**Core methods of `BoTorchModel`:** <br>\n",
"`fit` calls `Surrogate.fit` <br>\n",
"`predict` calls `Surrogate.predict` <br>\n",
"`gen` calls `Acquisition.optimize`"
]
},
{
"cell_type": "markdown",
"metadata": {
"code_folding": [],
"customInput": null,
"hidden_ranges": [],
"originalKey": "b6eaa943-3dbc-4f4a-b994-1e9d5d418bfa",
"showInput": false
},
"source": [
"# `BoTorchModel` instantiation"
]
},
{
"cell_type": "code",
"execution_count": 1,
"metadata": {
"code_folding": [],
"hidden_ranges": [],
"originalKey": "12d1d68e-b74b-455d-aeb9-ed9542db38c3"
},
"outputs": [],
"source": [
"from ax.models.torch.botorch_modular.acquisition import Acquisition\n",
"from ax.models.torch.botorch_modular.kg import KnowledgeGradient\n",
"from ax.models.torch.botorch_modular.model import BoTorchModel\n",
"from ax.models.torch.botorch_modular.surrogate import Surrogate\n",
"from botorch.models.gp_regression import FixedNoiseGP, SingleTaskGP\n",
"\n",
"# Explicit instantiation of `BoTorchModel`.\n",
"model = BoTorchModel(\n",
" surrogate=Surrogate(FixedNoiseGP),\n",
" acquisition_class=KnowledgeGradient, # This is a subclass of `Acquisition`.\n",
")"
]
},
{
"cell_type": "markdown",
"metadata": {
"code_folding": [],
"customInput": null,
"hidden_ranges": [],
"originalKey": "b915025a-8153-4f42-ae50-c3621dccb1b5",
"showInput": false
},
"source": [
"If `surrogate` and/or `acquisition_class` are not passed into the constructor, then they will auto-selected based on properties of the experiment, search space, and the data available for it [like so](https://github.com/facebook/Ax/blob/main/ax/models/torch/botorch_modular/utils.py)."
]
},
{
"cell_type": "code",
"execution_count": 2,
"metadata": {
"code_folding": [],
"hidden_ranges": [],
"originalKey": "034b0bac-e8ea-4519-9e03-54ffec9555f8"
},
"outputs": [],
"source": [
"# The surrogate is not specified, so it will be auto-selected during `model.fit`.\n",
"model = BoTorchModel(\n",
" acquisition_class=KnowledgeGradient\n",
")\n",
"\n",
"# The acquisition class is not specified, so it will be auto-selected during `model.gen`.\n",
"model = BoTorchModel(\n",
" surrogate=Surrogate(FixedNoiseGP)\n",
")\n",
"\n",
"# Both the surrogate and acquisition class will be auto-selected.\n",
"model = BoTorchModel()"
]
},
{
"cell_type": "markdown",
"metadata": {
"code_folding": [],
"customInput": null,
"hidden_ranges": [],
"originalKey": "581a5223-211f-4feb-abeb-57426cd55019",
"showInput": false
},
"source": [
"To use `ExpectedImprovement` and `NoisyExpectedImprovement`, initialize the `BoTorchModel` with the kwarg `botorch_acqf_class` instead of `acquisition_class`. By default, `acquisition_class` will be set to the base Ax `Acquisition` class."
]
},
{
"cell_type": "code",
"execution_count": 3,
"metadata": {
"code_folding": [],
"hidden_ranges": [],
"originalKey": "9ff9184e-50bf-404a-9b12-fa1ee9bd659c"
},
"outputs": [],
"source": [
"from botorch.acquisition.monte_carlo import qExpectedImprovement\n",
"from botorch.acquisition.monte_carlo import qNoisyExpectedImprovement\n",
"\n",
"EI_model = BoTorchModel(\n",
" surrogate=Surrogate(FixedNoiseGP),\n",
" botorch_acqf_class=qExpectedImprovement\n",
")\n",
"NEI_model = BoTorchModel(\n",
" surrogate=Surrogate(SingleTaskGP),\n",
" botorch_acqf_class=qNoisyExpectedImprovement\n",
")"
]
},
{
"cell_type": "markdown",
"metadata": {
"code_folding": [],
"customInput": null,
"hidden_ranges": [],
"originalKey": "0dfbc6b9-55bc-46de-8885-827e435a6f1a",
"showInput": false
},
"source": [
"# Use a Custom BoTorch `AcquisitionFunction`"
]
},
{
"cell_type": "markdown",
"metadata": {
"code_folding": [],
"customInput": null,
"hidden_ranges": [],
"originalKey": "b668732c-3a5b-48c3-b2e0-36fe9a61aefa",
"showInput": false
},
"source": [
"## Choose between the default Ax `Acquisition` class and creating a custom `Acquisition` subclass \n",
"In many cases, even when you want to use a custom BoTorch `AcquisitionFunction`, the default Ax `Acquisition` class may be enough **[Path 1]**. A custom Ax `Acquisition` subclass **[Path 2]** will be needed only if:\n",
"\n",
"- a custom acquisition function optimization method is required\n",
"- a custom “model dependency” is required, where a “model dependency” is defined as any value that is computed based on the state or properties of the `Surrogate` model and needs to be passed into the constructor for the `AcquisitionFunction`"
]
},
{
"cell_type": "markdown",
"metadata": {
"code_folding": [],
"customInput": null,
"hidden_ranges": [],
"originalKey": "ae275edd-8e58-4f97-af53-726248693e84",
"showInput": false
},
"source": [
"## [Path 1] Use the default Ax `Acquisition` class\n",
"\n",
"Construct your model in the same way that `ExpectedImprovement` and `NoisyExpectedImprovement` are constructed above. Then, if you want to set up storage for the model, skip to after **[Path 2]**."
]
},
{
"cell_type": "markdown",
"metadata": {
"code_folding": [],
"customInput": null,
"hidden_ranges": [],
"originalKey": "55a556f1-d073-4663-842a-975612679510",
"showInput": false
},
"source": [
"## [Path 2] Create a custom Ax `Acquisition` subclass\n",
"\n",
"To start, here is the inheritance tree for the `KnowledgeGradient` and `MultiFidelityKnowledgeGradient` subclasses. \n",
"\n",
"`Acquisition` <br>\n",
"↳ `MultiFidelityAcquisition(Acquisition)` <br>\n",
"↳ `KnowledgeGradient(Acquisition)` <br>\n",
" ↳ `MultiFidelityKnowledgeGradient(MultiFidelityAcquisition, KnowledgeGradient)` <br>\n",
"↳ `MyAcquisition(Acquisition)` **← your new subclass**\n",
"\n",
"The `Acquisition` class defines a default `optimize` function and `compute_model_dependencies` function. <br>\n",
" \n",
"**`optimize`**: makes a call to BoTorch's acquisition function optimizer with a specific set of kwargs specified by this function. <br>\n",
" \n",
"**`compute_model_dependencies`**: returns a dict of inputs to the BoTorch `AcquisitionFunction`.\n",
"\n",
"Creating a custom `Acquisition` subclass involves overriding either (or both) of these two functions."
]
},
{
"cell_type": "markdown",
"metadata": {
"code_folding": [],
"customInput": null,
"hidden_ranges": [],
"originalKey": "dd019da9-0796-43ea-83f4-9d10008ba191",
"showInput": false
},
"source": [
"First, create the structure for your `Acquisition` subclass. Each `Acquisition` subclass must have a BoTorch `AcquisitionFunction` class associated with it. We will add **`optimize`** and **`compute_model_dependencies`** to this class."
]
},
{
"cell_type": "code",
"execution_count": 4,
"metadata": {
"code_folding": [],
"hidden_ranges": [],
"originalKey": "dd0eea28-8558-4cde-a131-dceec090248c",
"showInput": true
},
"outputs": [],
"source": [
"# `qKnowledgeGradient` is being used as a placeholder here.\n",
"\n",
"# from botorch.acquisition.my_acquisition import qMyAcquisition\n",
"from botorch.acquisition.knowledge_gradient import qKnowledgeGradient\n",
"\n",
"class MyAcquisition(Acquisition):\n",
" # default_botorch_acqf_class = qMyAcquisition\n",
" default_botorch_acqf_class = qKnowledgeGradient\n"
]
},
{
"cell_type": "markdown",
"metadata": {
"code_folding": [],
"customInput": null,
"hidden_ranges": [],
"originalKey": "01a2e99b-fa8b-4d57-802c-1a47250d0ee0",
"showInput": false
},
"source": [
"### [Path 2: Step 1] Override `Acquisition.optimize`\n",
"\n",
"By default, the `Acquisition` subclasses run `super().optimize` but they specify their own `optimizer_options` (if needed). These `optimizer_options` are then sent into [BoTorch's `optimize_acqf` optimizer](https://github.com/pytorch/botorch/blob/main/botorch/optim/optimize.py).\n",
"\n",
"The following arguments are always passed into the optimizer:\n",
"- `bounds`\n",
"- `q`\n",
"- `inequality_constraints`\n",
"- `fixed_features`\n",
"- `post_processing_func`\n",
"\n",
"Any kwargs that `optimize_acqf` takes in that are not part of this list can be set by `optimizer_options`.\n",
"\n",
"**NOTE:** If the optimizer for the BoTorch `AcquisitionFunction` that you want to use does not require any kwargs other than those in the list, then you do not need to override `Acquisition.optimize`."
]
},
{
"cell_type": "markdown",
"metadata": {
"code_folding": [],
"customInput": null,
"hidden_ranges": [],
"originalKey": "7d4cd54e-8766-4d89-b8a8-15b778eb46ff",
"showInput": false
},
"source": [
"As an example, for `MaxValueEntropySearch`, we want to use \"sequential greedy\" optimization of the acquisition function with a batch of `q > 1` candidates. So, we want `sequential=True` to be passed into `optimize_acqf`:"
]
},
{
"cell_type": "code",
"execution_count": 5,
"metadata": {
"code_folding": [],
"hidden_ranges": [],
"originalKey": "52bcaf22-7c44-43dc-8377-cd7275c719ed",
"showInput": true
},
"outputs": [],
"source": [
"from typing import Any, Callable, Dict, List, Optional, Tuple\n",
"from ax.models.torch.botorch_modular.acquisition import Optimizer\n",
"from torch import Tensor\n",
"\n",
"def optimize(\n",
" self,\n",
" bounds: Tensor,\n",
" n: int,\n",
" optimizer_class: Optional[Optimizer] = None,\n",
" inequality_constraints: Optional[List[Tuple[Tensor, Tensor, float]]] = None,\n",
" fixed_features: Optional[Dict[int, float]] = None,\n",
" rounding_func: Optional[Callable[[Tensor], Tensor]] = None,\n",
" optimizer_options: Optional[Dict[str, Any]] = None,\n",
") -> Tuple[Tensor, Tensor]:\n",
" optimizer_options = optimizer_options or {}\n",
" optimizer_options[\"sequential\"] = True\n",
" return super().optimize(\n",
" bounds=bounds,\n",
" n=n,\n",
" inequality_constraints=None,\n",
" fixed_features=fixed_features,\n",
" rounding_func=rounding_func,\n",
" optimizer_options=optimizer_options,\n",
" )"
]
},
{
"cell_type": "markdown",
"metadata": {
"code_folding": [],
"customInput": null,
"hidden_ranges": [],
"originalKey": "953a43e7-72c7-4aa6-b39d-8512b0e29aff",
"showInput": false
},
"source": [
"And with that, we are done overriding `Acquisition.optimize`."
]
},
{
"cell_type": "markdown",
"metadata": {
"code_folding": [],
"customInput": null,
"hidden_ranges": [],
"originalKey": "28930608-cf7f-491c-a1f7-e32387433ff3",
"showInput": false
},
"source": [
"### [Path 2: Step 2] Override Acquisition.compute_model_dependencies\n",
"\n",
"Similar to the base `Acquisition.optimize`, the `Acquisition` subclasses run `super().compute_model_dependencies` but they add to the dictionary their own dependencies (if any). This `model_deps` dictionary is then sent into the BoTorch `AcquisitionFunction` constructor as `**model_deps`.\n",
"\n",
"**NOTE:** If the BoTorch `AcquisitionFunction` that you want to use does not require any special `__init__` arguments other than `model`, `objective`, `X_pending`, and `X_baseline`, then you do not need to override `Acquisition.compute_model_dependencies`."
]
},
{
"cell_type": "markdown",
"metadata": {
"code_folding": [],
"customInput": null,
"hidden_ranges": [],
"originalKey": "005dbad9-2eec-446f-9ab3-2fbf84051370",
"showInput": false
},
"source": [
"As an example, for `MaxValueEntropySearch`, we must pass into `qMaxValueEntropy.__init__` a `candidate_set: Tensor` and we also want to specify `maximize: bool`. For the exact code, [see here](https://github.com/facebook/Ax/blob/main/ax/models/torch/botorch_modular/mes.py)."
]
},
{
"cell_type": "code",
"execution_count": 6,
"metadata": {
"code_folding": [],
"hidden_ranges": [],
"originalKey": "fef17722-fb2b-4cf7-b5bd-4adbcbcaa054",
"showInput": true
},
"outputs": [],
"source": [
"@classmethod\n",
"def compute_model_dependencies(\n",
" cls,\n",
" surrogate: Surrogate,\n",
" bounds: List[Tuple[float, float]],\n",
" objective_weights: Tensor,\n",
" pending_observations: Optional[List[Tensor]] = None,\n",
" outcome_constraints: Optional[Tuple[Tensor, Tensor]] = None,\n",
" linear_constraints: Optional[Tuple[Tensor, Tensor]] = None,\n",
" fixed_features: Optional[Dict[int, float]] = None,\n",
" target_fidelities: Optional[Dict[int, float]] = None,\n",
" options: Optional[Dict[str, Any]] = None,\n",
") -> Dict[str, Any]:\n",
"\n",
" # Get the dependencies of the parent class.\n",
" dependencies = super().compute_model_dependencies(...)\n",
"\n",
" # Calculate `candidate_set`.\n",
" candidate_set = ...\n",
" # Calculate `maximize`.\n",
" maximize = ...\n",
"\n",
" # Update and return the model dependencies.\n",
" dependencies.update(\n",
" {\"candidate_set\": candidate_set, \"maximize\": maximize}\n",
" )\n",
" return dependencies"
]
},
{
"cell_type": "markdown",
"metadata": {
"code_folding": [],
"customInput": null,
"hidden_ranges": [],
"originalKey": "169feed0-9ee7-40c6-b6c2-023b8eb11c8e",
"showInput": false
},
"source": [
"And with that, we are done overriding `Acquisition.compute_model_dependencies`."
]
},
{
"cell_type": "markdown",
"metadata": {
"code_folding": [],
"customInput": null,
"hidden_ranges": [],
"originalKey": "5ffc51b3-2256-47e1-9d20-4283221a411e",
"showInput": false
},
"source": [
"### [Path 2: Step 3] Put it all together and try it out"
]
},
{
"cell_type": "code",
"execution_count": 7,
"metadata": {
"code_folding": [],
"hidden_ranges": [],
"originalKey": "99050062-63b1-4c29-a0ff-58742e660c96",
"showInput": true
},
"outputs": [],
"source": [
"# `qKnowledgeGradient` is being used as a placeholder here.\n",
"\n",
"# from botorch.acquisition.my_acquisition import qMyAcquisition\n",
"from botorch.acquisition.knowledge_gradient import qKnowledgeGradient\n",
"\n",
"class MyAcquisition(Acquisition):\n",
" # default_botorch_acqf_class = qMyAcquisition\n",
" default_botorch_acqf_class = qKnowledgeGradient\n",
" \n",
" def optimize(\n",
" self,\n",
" # other kwargs,\n",
" ) -> Tuple[Tensor, Tensor]:\n",
" ...\n",
" return super().optimize(...)\n",
" \n",
" @classmethod\n",
" def compute_model_dependencies(\n",
" cls,\n",
" # other kwargs,\n",
" ) -> Dict[str, Any]:\n",
" dependencies = super().compute_model_dependencies(...)\n",
" ...\n",
" return dependencies"
]
},
{
"cell_type": "markdown",
"metadata": {
"code_folding": [],
"customInput": null,
"hidden_ranges": [],
"originalKey": "368af6d0-4d8a-4b9d-9873-9bf093fba451",
"showInput": false
},
"source": [
"Now, we can use the new custom `Acquisition` subclass in the same way as we do for `KnowledgeGradient` in the **`BoTorchModel` instantiation** section."
]
},
{
"cell_type": "markdown",
"metadata": {
"code_folding": [],
"customInput": null,
"hidden_ranges": [],
"originalKey": "a201541d-2670-46d4-9358-2fa53d3fe55d",
"showInput": false
},
"source": [
"## Set up storage for the new setup\n",
"\n",
"Optionally, to allow the Ax models to be serializable (and allow for resumable optimization via JSON or SQL storage), navigate [here](https://github.com/facebook/Ax/blob/main/ax/storage/botorch_modular_registry.py).\n",
"\n",
"1. If you created a new `Acquisition` subclass, add it to `ACQUISITION_REGISTRY`.\n",
"2. Add the corresponding BoTorch `AcquisitionFunction` class to `ACQUISITION_FUNCTION_REGISTRY`."
]
},
{
"cell_type": "markdown",
"metadata": {
"code_folding": [],
"customInput": null,
"hidden_ranges": [],
"originalKey": "055bce82-fcef-4ed8-b251-605250c5b947",
"showInput": false
},
"source": [
"```\n",
"ACQUISITION_REGISTRY: Dict[Type[Acquisition], int] = {\n",
" Acquisition: 0,\n",
" KnowledgeGradient: 1,\n",
" ...\n",
" MyAcquisition: 5, # Add this line.\n",
"}\n",
"\n",
"ACQUISITION_FUNCTION_REGISTRY: Dict[Type[AcquisitionFunction], int] = {\n",
" qExpectedImprovement: 0,\n",
" qNoisyExpectedImprovement: 1,\n",
" ...\n",
" qMyAcquisition: 6, # Add this line.\n",
"}\n",
"```"
]
},
{
"cell_type": "markdown",
"metadata": {
"code_folding": [],
"customInput": null,
"hidden_ranges": [],
"originalKey": "6ee7ee9d-b5ad-46a5-bcef-228892ff94ca",
"showInput": false
},
"source": [
"# Details of `BoTorchModel` Subcomponent Classes"
]
},
{
"cell_type": "markdown",
"metadata": {
"code_folding": [],
"customInput": null,
"hidden_ranges": [],
"originalKey": "d652d0e5-45ca-48a9-ad95-a275f05e722d",
"showInput": false
},
"source": [
"## class `Surrogate`\n",
"\n",
"Ax wrapper for BoTorch GP `Model` classes. Optionally, a BoTorch `MarginalLogLikelihood` class can also be passed into the construction of Surrogate.\n",
"\n",
"**Core methods:** `fit`, `predict`, `update`, `construct`, `best_in_sample_point`, `best_out_of_sample_point`\n",
"These core methods are all called by BoTorchModel internally.\n",
"\n",
"\n",
"\n",
"\n",
"\n"
]
},
{
"cell_type": "code",
"execution_count": 8,
"metadata": {
"originalKey": "994b3976-2486-4a4b-9c34-4c2ba93fe816"
},
"outputs": [],
"source": [
"from botorch.models.gp_regression import FixedNoiseGP\n",
"from gpytorch.mlls.exact_marginal_log_likelihood import ExactMarginalLogLikelihood\n",
"\n",
"surrogate = Surrogate(\n",
" botorch_model_class=FixedNoiseGP, # required kwarg\n",
" mll_class=ExactMarginalLogLikelihood, # optional kwarg\n",
")"
]
},
{
"cell_type": "markdown",
"metadata": {
"code_folding": [],
"customInput": null,
"hidden_ranges": [],
"originalKey": "acdc8efc-5376-46b3-b0d4-3976cdb6a828",
"showInput": false
},
"source": [
"## class `Acquisition`\n",
"\n",
"Base Ax class for BoTorch `AcquisitionFunction` classes.\n",
"\n",
"**Core methods:** `optimize`, `compute_model_dependencies`"
]
}
],
"metadata": {
"kernelspec": {
"display_name": "python3",
"language": "python",
"name": "python3"
}
},
"nbformat": 4,
"nbformat_minor": 2
}