From 4fcaa91f10fc2c237c21004a3f96dd858db43ba1 Mon Sep 17 00:00:00 2001 From: Bernie Beckerman Date: Tue, 4 Jun 2024 10:59:56 -0700 Subject: [PATCH] Tutorial to mimic parameter constraint functionality for nonlinear constraints (#2500) Summary: This diff adds to tutorials a method to early-exit trial evaluation based on an easy-to-calculate constraint metric (e.g., a nonlinear function of parameter values) in order to mimic behavior of parameter constraints when parameter constraints cannot be used. See this discussion for details ([link](https://github.com/facebook/Ax/issues/2460?fbclid=IwAR1rpJpyOP37z32ACczQRcWc51TQF640RLr17DYQU7sCyOLvuWOim7UCGIY#issuecomment-2115746662)). Adds the following to the end of the `Special Cases` section of the Service API tutorial ([link](https://ax.dev/tutorials/gpei_hartmann_service.html#Special-Cases)): {F1670348258} Differential Revision: D58146231 --- tutorials/gpei_hartmann_service.ipynb | 1004 +++++++++++++------------ 1 file changed, 523 insertions(+), 481 deletions(-) diff --git a/tutorials/gpei_hartmann_service.ipynb b/tutorials/gpei_hartmann_service.ipynb index b6c225fccd7..177f7a182d4 100644 --- a/tutorials/gpei_hartmann_service.ipynb +++ b/tutorials/gpei_hartmann_service.ipynb @@ -1,483 +1,525 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "# Service API Example on Hartmann6\n", - "\n", - "The Ax Service API is designed to allow the user to control scheduling of trials and data computation while having an easy to use interface with Ax.\n", - "\n", - "The user iteratively:\n", - "- Queries Ax for candidates\n", - "- Schedules / deploys them however they choose\n", - "- Computes data and logs to Ax\n", - "- Repeat" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "from ax.service.ax_client import AxClient, ObjectiveProperties\n", - "from ax.utils.measurement.synthetic_functions import hartmann6\n", - "from ax.utils.notebook.plotting import init_notebook_plotting, render\n", - "\n", - "init_notebook_plotting()" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "## 1. Initialize client\n", - "\n", - "Create a client object to interface with Ax APIs. By default this runs locally without storage." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "ax_client = AxClient()" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "## 2. Set up experiment\n", - "An experiment consists of a **search space** (parameters and parameter constraints) and **optimization configuration** (objectives and outcome constraints). Note that:\n", - "- Only `parameters`, and `objectives` arguments are required.\n", - "- Dictionaries in `parameters` have the following required keys: \"name\" - parameter name, \"type\" - parameter type (\"range\", \"choice\" or \"fixed\"), \"bounds\" for range parameters, \"values\" for choice parameters, and \"value\" for fixed parameters.\n", - "- Dictionaries in `parameters` can optionally include \"value_type\" (\"int\", \"float\", \"bool\" or \"str\"), \"log_scale\" flag for range parameters, and \"is_ordered\" flag for choice parameters.\n", - "- `parameter_constraints` should be a list of strings of form \"p1 >= p2\" or \"p1 + p2 <= some_bound\".\n", - "- `outcome_constraints` should be a list of strings of form \"constrained_metric <= some_bound\"." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "ax_client.create_experiment(\n", - " name=\"hartmann_test_experiment\",\n", - " parameters=[\n", - " {\n", - " \"name\": \"x1\",\n", - " \"type\": \"range\",\n", - " \"bounds\": [0.0, 1.0],\n", - " \"value_type\": \"float\", # Optional, defaults to inference from type of \"bounds\".\n", - " \"log_scale\": False, # Optional, defaults to False.\n", - " },\n", - " {\n", - " \"name\": \"x2\",\n", - " \"type\": \"range\",\n", - " \"bounds\": [0.0, 1.0],\n", - " },\n", - " {\n", - " \"name\": \"x3\",\n", - " \"type\": \"range\",\n", - " \"bounds\": [0.0, 1.0],\n", - " },\n", - " {\n", - " \"name\": \"x4\",\n", - " \"type\": \"range\",\n", - " \"bounds\": [0.0, 1.0],\n", - " },\n", - " {\n", - " \"name\": \"x5\",\n", - " \"type\": \"range\",\n", - " \"bounds\": [0.0, 1.0],\n", - " },\n", - " {\n", - " \"name\": \"x6\",\n", - " \"type\": \"range\",\n", - " \"bounds\": [0.0, 1.0],\n", - " },\n", - " ],\n", - " objectives={\"hartmann6\": ObjectiveProperties(minimize=True)},\n", - " parameter_constraints=[\"x1 + x2 <= 2.0\"], # Optional.\n", - " outcome_constraints=[\"l2norm <= 1.25\"], # Optional.\n", - ")" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "## 3. Define how to evaluate trials\n", - "When using Ax a service, evaluation of parameterizations suggested by Ax is done either locally or, more commonly, using an external scheduler. Below is a dummy evaluation function that outputs data for two metrics \"hartmann6\" and \"l2norm\". Note that all returned metrics correspond to either the `objectives` set on experiment creation or the metric names mentioned in `outcome_constraints`." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "import numpy as np\n", - "\n", - "\n", - "def evaluate(parameterization):\n", - " x = np.array([parameterization.get(f\"x{i+1}\") for i in range(6)])\n", - " # In our case, standard error is 0, since we are computing a synthetic function.\n", - " return {\"hartmann6\": (hartmann6(x), 0.0), \"l2norm\": (np.sqrt((x**2).sum()), 0.0)}" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Result of the evaluation should generally be a mapping of the format: `{metric_name -> (mean, SEM)}`. If there is only one metric in the experiment – the objective – then evaluation function can return a single tuple of mean and SEM, in which case Ax will assume that evaluation corresponds to the objective. _It can also return only the mean as a float, in which case Ax will treat SEM as unknown and use a model that can infer it._ \n", - "\n", - "For more details on evaluation function, refer to the \"Trial Evaluation\" section in the Ax docs at [ax.dev](https://ax.dev/)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "## 4. Run optimization loop\n", - "With the experiment set up, we can start the optimization loop.\n", - "\n", - "At each step, the user queries the client for a new trial then submits the evaluation of that trial back to the client.\n", - "\n", - "Note that Ax auto-selects an appropriate optimization algorithm based on the search space. For more advance use cases that require a specific optimization algorithm, pass a `generation_strategy` argument into the `AxClient` constructor. Note that when Bayesian Optimization is used, generating new trials may take a few minutes." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "for i in range(25):\n", - " parameterization, trial_index = ax_client.get_next_trial()\n", - " # Local evaluation here can be replaced with deployment to external system.\n", - " ax_client.complete_trial(trial_index=trial_index, raw_data=evaluate(parameterization))" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "### How many trials can run in parallel?\n", - "By default, Ax restricts number of trials that can run in parallel for some optimization stages, in order to improve the optimization performance and reduce the number of trials that the optimization will require. To check the maximum parallelism for each optimization stage:" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "ax_client.get_max_parallelism()" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The output of this function is a list of tuples of form (number of trials, max parallelism), so the example above means \"the max parallelism is 12 for the first 12 trials and 3 for all subsequent trials.\" This is because the first 12 trials are produced quasi-randomly and can all be evaluated at once, and subsequent trials are produced via Bayesian optimization, which converges on optimal point in fewer trials when parallelism is limited. `MaxParallelismReachedException` indicates that the parallelism limit has been reached –– refer to the 'Service API Exceptions Meaning and Handling' section at the end of the tutorial for handling." - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "### How to view all existing trials during optimization?" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "ax_client.generation_strategy.trials_as_df" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "## 5. Retrieve best parameters\n", - "\n", - "Once it's complete, we can access the best parameters found, as well as the corresponding metric values." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "best_parameters, values = ax_client.get_best_parameters()\n", - "best_parameters" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "means, covariances = values\n", - "means" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "For comparison, Hartmann6 minimum:" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "hartmann6.fmin" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "## 6. Plot the response surface and optimization trace\n", - "Here we arbitrarily select \"x1\" and \"x2\" as the two parameters to plot for both metrics, \"hartmann6\" and \"l2norm\"." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "render(ax_client.get_contour_plot())" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "We can also retrieve a contour plot for the other metric, \"l2norm\" –– say, we are interested in seeing the response surface for parameters \"x3\" and \"x4\" for this one." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "render(ax_client.get_contour_plot(param_x=\"x3\", param_y=\"x4\", metric_name=\"l2norm\"))" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Here we plot the optimization trace, showing the progression of finding the point with the optimal objective:" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "render(\n", - " ax_client.get_optimization_trace(objective_optimum=hartmann6.fmin)\n", - ") # Objective_optimum is optional." - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "## 7. Save / reload optimization to JSON / SQL\n", - "We can serialize the state of optimization to JSON and save it to a `.json` file or save it to the SQL backend. For the former:" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "ax_client.save_to_json_file() # For custom filepath, pass `filepath` argument." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "restored_ax_client = (\n", - " AxClient.load_from_json_file()\n", - ") # For custom filepath, pass `filepath` argument." - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "To store state of optimization to an SQL backend, first follow [setup instructions](https://ax.dev/docs/storage.html#sql) on Ax website." - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Having set up the SQL backend, pass `DBSettings` to `AxClient` on instantiation (note that `SQLAlchemy` dependency will have to be installed – for installation, refer to [optional dependencies](https://ax.dev/docs/installation.html#optional-dependencies) on Ax website):" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "from ax.storage.sqa_store.structs import DBSettings\n", - "\n", - "# URL is of the form \"dialect+driver://username:password@host:port/database\".\n", - "db_settings = DBSettings(url=\"sqlite:///foo.db\")\n", - "# Instead of URL, can provide a `creator function`; can specify custom encoders/decoders if necessary.\n", - "new_ax = AxClient(db_settings=db_settings)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "When valid `DBSettings` are passed into `AxClient`, a unique experiment name is a required argument (`name`) to `ax_client.create_experiment`. The **state of the optimization is auto-saved** any time it changes (i.e. a new trial is added or completed, etc). \n", - "\n", - "To reload an optimization state later, instantiate `AxClient` with the same `DBSettings` and use `ax_client.load_experiment_from_database(experiment_name=\"my_experiment\")`." - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "# Special Cases" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "**Evaluation failure**: should any optimization iterations fail during evaluation, `log_trial_failure` will ensure that the same trial is not proposed again." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "_, trial_index = ax_client.get_next_trial()\n", - "ax_client.log_trial_failure(trial_index=trial_index)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "**Adding custom trials**: should there be need to evaluate a specific parameterization, `attach_trial` will add it to the experiment." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "ax_client.attach_trial(\n", - " parameters={\"x1\": 0.9, \"x2\": 0.9, \"x3\": 0.9, \"x4\": 0.9, \"x5\": 0.9, \"x6\": 0.9}\n", - ")" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "**Need to run many trials in parallel**: for optimal results and optimization efficiency, we strongly recommend sequential optimization (generating a few trials, then waiting for them to be completed with evaluation data). However, if your use case needs to dispatch many trials in parallel before they are updated with data and you are running into the *\"All trials for current model have been generated, but not enough data has been observed to fit next model\"* error, instantiate `AxClient` as `AxClient(enforce_sequential_optimization=False)`." - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "# Service API Exceptions Meaning and Handling\n", - "[**`DataRequiredError`**](https://ax.dev/api/exceptions.html#ax.exceptions.core.DataRequiredError): Ax generation strategy needs to be updated with more data to proceed to the next optimization model. When the optimization moves from initialization stage to the Bayesian optimization stage, the underlying BayesOpt model needs sufficient data to train. For optimal results and optimization efficiency (finding the optimal point in the least number of trials), we recommend sequential optimization (generating a few trials, then waiting for them to be completed with evaluation data). Therefore, the correct way to handle this exception is to wait until more trial evaluations complete and log their data via `ax_client.complete_trial(...)`. \n", - "\n", - "However, if there is strong need to generate more trials before more data is available, instantiate `AxClient` as `AxClient(enforce_sequential_optimization=False)`. With this setting, as many trials will be generated from the initialization stage as requested, and the optimization will move to the BayesOpt stage whenever enough trials are completed." - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "[**`MaxParallelismReachedException`**](https://ax.dev/api/modelbridge.html#ax.modelbridge.generation_strategy.MaxParallelismReachedException): generation strategy restricts the number of trials that can be ran simultaneously (to encourage sequential optimization), and the parallelism limit has been reached. The correct way to handle this exception is the same as `DataRequiredError` – to wait until more trial evluations complete and log their data via `ax_client.complete_trial(...)`.\n", - " \n", - "In some cases higher parallelism is important, so `enforce_sequential_optimization=False` kwarg to AxClient allows to suppress limiting of parallelism. It's also possible to override the default parallelism setting for all stages of the optimization by passing `choose_generation_strategy_kwargs` to `ax_client.create_experiment`:" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "ax_client = AxClient()\n", - "ax_client.create_experiment(\n", - " parameters=[\n", - " {\"name\": \"x\", \"type\": \"range\", \"bounds\": [-5.0, 10.0]},\n", - " {\"name\": \"y\", \"type\": \"range\", \"bounds\": [0.0, 15.0]},\n", - " ],\n", - " # Sets max parallelism to 10 for all steps of the generation strategy.\n", - " choose_generation_strategy_kwargs={\"max_parallelism_override\": 10},\n", - ")" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "ax_client.get_max_parallelism() # Max parallelism is now 10 for all stages of the optimization." - ] - } - ], - "metadata": { - "kernelspec": { - "display_name": "Python 3 (ipykernel)", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.11.5" - } - }, - "nbformat": 4, - "nbformat_minor": 2 + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# Service API Example on Hartmann6\n", + "\n", + "The Ax Service API is designed to allow the user to control scheduling of trials and data computation while having an easy to use interface with Ax.\n", + "\n", + "The user iteratively:\n", + "- Queries Ax for candidates\n", + "- Schedules / deploys them however they choose\n", + "- Computes data and logs to Ax\n", + "- Repeat" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "from ax.service.ax_client import AxClient, ObjectiveProperties\n", + "from ax.utils.measurement.synthetic_functions import hartmann6\n", + "from ax.utils.notebook.plotting import init_notebook_plotting, render\n", + "\n", + "init_notebook_plotting()" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 1. Initialize client\n", + "\n", + "Create a client object to interface with Ax APIs. By default this runs locally without storage." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "ax_client = AxClient()" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 2. Set up experiment\n", + "An experiment consists of a **search space** (parameters and parameter constraints) and **optimization configuration** (objectives and outcome constraints). Note that:\n", + "- Only `parameters`, and `objectives` arguments are required.\n", + "- Dictionaries in `parameters` have the following required keys: \"name\" - parameter name, \"type\" - parameter type (\"range\", \"choice\" or \"fixed\"), \"bounds\" for range parameters, \"values\" for choice parameters, and \"value\" for fixed parameters.\n", + "- Dictionaries in `parameters` can optionally include \"value_type\" (\"int\", \"float\", \"bool\" or \"str\"), \"log_scale\" flag for range parameters, and \"is_ordered\" flag for choice parameters.\n", + "- `parameter_constraints` should be a list of strings of form \"p1 >= p2\" or \"p1 + p2 <= some_bound\".\n", + "- `outcome_constraints` should be a list of strings of form \"constrained_metric <= some_bound\"." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "ax_client.create_experiment(\n", + " name=\"hartmann_test_experiment\",\n", + " parameters=[\n", + " {\n", + " \"name\": \"x1\",\n", + " \"type\": \"range\",\n", + " \"bounds\": [0.0, 1.0],\n", + " \"value_type\": \"float\", # Optional, defaults to inference from type of \"bounds\".\n", + " \"log_scale\": False, # Optional, defaults to False.\n", + " },\n", + " {\n", + " \"name\": \"x2\",\n", + " \"type\": \"range\",\n", + " \"bounds\": [0.0, 1.0],\n", + " },\n", + " {\n", + " \"name\": \"x3\",\n", + " \"type\": \"range\",\n", + " \"bounds\": [0.0, 1.0],\n", + " },\n", + " {\n", + " \"name\": \"x4\",\n", + " \"type\": \"range\",\n", + " \"bounds\": [0.0, 1.0],\n", + " },\n", + " {\n", + " \"name\": \"x5\",\n", + " \"type\": \"range\",\n", + " \"bounds\": [0.0, 1.0],\n", + " },\n", + " {\n", + " \"name\": \"x6\",\n", + " \"type\": \"range\",\n", + " \"bounds\": [0.0, 1.0],\n", + " },\n", + " ],\n", + " objectives={\"hartmann6\": ObjectiveProperties(minimize=True)},\n", + " parameter_constraints=[\"x1 + x2 <= 2.0\"], # Optional.\n", + " outcome_constraints=[\"l2norm <= 1.25\"], # Optional.\n", + ")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 3. Define how to evaluate trials\n", + "When using Ax a service, evaluation of parameterizations suggested by Ax is done either locally or, more commonly, using an external scheduler. Below is a dummy evaluation function that outputs data for two metrics \"hartmann6\" and \"l2norm\". Note that all returned metrics correspond to either the `objectives` set on experiment creation or the metric names mentioned in `outcome_constraints`." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "import numpy as np\n", + "\n", + "\n", + "def evaluate(parameterization):\n", + " x = np.array([parameterization.get(f\"x{i+1}\") for i in range(6)])\n", + " # In our case, standard error is 0, since we are computing a synthetic function.\n", + " return {\"hartmann6\": (hartmann6(x), 0.0), \"l2norm\": (np.sqrt((x**2).sum()), 0.0)}" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Result of the evaluation should generally be a mapping of the format: `{metric_name -> (mean, SEM)}`. If there is only one metric in the experiment – the objective – then evaluation function can return a single tuple of mean and SEM, in which case Ax will assume that evaluation corresponds to the objective. _It can also return only the mean as a float, in which case Ax will treat SEM as unknown and use a model that can infer it._ \n", + "\n", + "For more details on evaluation function, refer to the \"Trial Evaluation\" section in the Ax docs at [ax.dev](https://ax.dev/)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 4. Run optimization loop\n", + "With the experiment set up, we can start the optimization loop.\n", + "\n", + "At each step, the user queries the client for a new trial then submits the evaluation of that trial back to the client.\n", + "\n", + "Note that Ax auto-selects an appropriate optimization algorithm based on the search space. For more advance use cases that require a specific optimization algorithm, pass a `generation_strategy` argument into the `AxClient` constructor. Note that when Bayesian Optimization is used, generating new trials may take a few minutes." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "for i in range(25):\n", + " parameterization, trial_index = ax_client.get_next_trial()\n", + " # Local evaluation here can be replaced with deployment to external system.\n", + " ax_client.complete_trial(trial_index=trial_index, raw_data=evaluate(parameterization))" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### How many trials can run in parallel?\n", + "By default, Ax restricts number of trials that can run in parallel for some optimization stages, in order to improve the optimization performance and reduce the number of trials that the optimization will require. To check the maximum parallelism for each optimization stage:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "ax_client.get_max_parallelism()" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The output of this function is a list of tuples of form (number of trials, max parallelism), so the example above means \"the max parallelism is 12 for the first 12 trials and 3 for all subsequent trials.\" This is because the first 12 trials are produced quasi-randomly and can all be evaluated at once, and subsequent trials are produced via Bayesian optimization, which converges on optimal point in fewer trials when parallelism is limited. `MaxParallelismReachedException` indicates that the parallelism limit has been reached –– refer to the 'Service API Exceptions Meaning and Handling' section at the end of the tutorial for handling." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### How to view all existing trials during optimization?" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "ax_client.generation_strategy.trials_as_df" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 5. Retrieve best parameters\n", + "\n", + "Once it's complete, we can access the best parameters found, as well as the corresponding metric values." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "best_parameters, values = ax_client.get_best_parameters()\n", + "best_parameters" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "means, covariances = values\n", + "means" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "For comparison, Hartmann6 minimum:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "hartmann6.fmin" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 6. Plot the response surface and optimization trace\n", + "Here we arbitrarily select \"x1\" and \"x2\" as the two parameters to plot for both metrics, \"hartmann6\" and \"l2norm\"." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "render(ax_client.get_contour_plot())" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "We can also retrieve a contour plot for the other metric, \"l2norm\" –– say, we are interested in seeing the response surface for parameters \"x3\" and \"x4\" for this one." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "render(ax_client.get_contour_plot(param_x=\"x3\", param_y=\"x4\", metric_name=\"l2norm\"))" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Here we plot the optimization trace, showing the progression of finding the point with the optimal objective:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "render(\n", + " ax_client.get_optimization_trace(objective_optimum=hartmann6.fmin)\n", + ") # Objective_optimum is optional." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 7. Save / reload optimization to JSON / SQL\n", + "We can serialize the state of optimization to JSON and save it to a `.json` file or save it to the SQL backend. For the former:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "ax_client.save_to_json_file() # For custom filepath, pass `filepath` argument." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "restored_ax_client = (\n", + " AxClient.load_from_json_file()\n", + ") # For custom filepath, pass `filepath` argument." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "To store state of optimization to an SQL backend, first follow [setup instructions](https://ax.dev/docs/storage.html#sql) on Ax website." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Having set up the SQL backend, pass `DBSettings` to `AxClient` on instantiation (note that `SQLAlchemy` dependency will have to be installed – for installation, refer to [optional dependencies](https://ax.dev/docs/installation.html#optional-dependencies) on Ax website):" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "from ax.storage.sqa_store.structs import DBSettings\n", + "\n", + "# URL is of the form \"dialect+driver://username:password@host:port/database\".\n", + "db_settings = DBSettings(url=\"sqlite:///foo.db\")\n", + "# Instead of URL, can provide a `creator function`; can specify custom encoders/decoders if necessary.\n", + "new_ax = AxClient(db_settings=db_settings)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "When valid `DBSettings` are passed into `AxClient`, a unique experiment name is a required argument (`name`) to `ax_client.create_experiment`. The **state of the optimization is auto-saved** any time it changes (i.e. a new trial is added or completed, etc). \n", + "\n", + "To reload an optimization state later, instantiate `AxClient` with the same `DBSettings` and use `ax_client.load_experiment_from_database(experiment_name=\"my_experiment\")`." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# Special Cases" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "**Evaluation failure**: should any optimization iterations fail during evaluation, `log_trial_failure` will ensure that the same trial is not proposed again." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "_, trial_index = ax_client.get_next_trial()\n", + "ax_client.log_trial_failure(trial_index=trial_index)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "**Adding custom trials**: should there be need to evaluate a specific parameterization, `attach_trial` will add it to the experiment." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "ax_client.attach_trial(\n", + " parameters={\"x1\": 0.9, \"x2\": 0.9, \"x3\": 0.9, \"x4\": 0.9, \"x5\": 0.9, \"x6\": 0.9}\n", + ")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "**Need to run many trials in parallel**: for optimal results and optimization efficiency, we strongly recommend sequential optimization (generating a few trials, then waiting for them to be completed with evaluation data). However, if your use case needs to dispatch many trials in parallel before they are updated with data and you are running into the *\"All trials for current model have been generated, but not enough data has been observed to fit next model\"* error, instantiate `AxClient` as `AxClient(enforce_sequential_optimization=False)`." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "**Nonlinear parameter constraints and/or constraints on non-Range parameters:** Ax parameter constraints can currently only support linear inequalities ([discussion](https://github.com/facebook/Ax/issues/153)). Users may be able to simulate this functionality, however, by substituting the following `evaluate` function for that defined in section 3 above." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "def evaluate(parameterization):\n", + " x = np.array([parameterization.get(f\"x{i+1}\") for i in range(6)])\n", + " # First calculate the nonlinear quantity to be constrained.\n", + " l2norm = np.sqrt((x**2).sum())\n", + " # Then define a constraint consistent with an outcome constraint on this experiment.\n", + " if l2norm > 1.25:\n", + " return {\"l2norm\": (l2norm, 0.0)}\n", + " return {\"hartmann6\": (hartmann6(x), 0.0), \"l2norm\": (l2norm, 0.0)}" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "For this to work, the constraint quantity (`l2norm` in this case) should have a corresponding outcome constraint on the experiment. See the outcome_constraint arg to ax_client.create_experiment in section 2 above for how to specify outcome constraints.\n", + "\n", + "This setup accomplishes the following:\n", + "1. Allows computation of an arbitrarily complex constraint value.\n", + "2. Skips objective computation when the constraint is violated, useful when the objective is relatively expensive to compute.\n", + "3. Constraint metric values are returned even when there is a violation. This helps the model understand + avoid constraint violations." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# Service API Exceptions Meaning and Handling\n", + "[**`DataRequiredError`**](https://ax.dev/api/exceptions.html#ax.exceptions.core.DataRequiredError): Ax generation strategy needs to be updated with more data to proceed to the next optimization model. When the optimization moves from initialization stage to the Bayesian optimization stage, the underlying BayesOpt model needs sufficient data to train. For optimal results and optimization efficiency (finding the optimal point in the least number of trials), we recommend sequential optimization (generating a few trials, then waiting for them to be completed with evaluation data). Therefore, the correct way to handle this exception is to wait until more trial evaluations complete and log their data via `ax_client.complete_trial(...)`. \n", + "\n", + "However, if there is strong need to generate more trials before more data is available, instantiate `AxClient` as `AxClient(enforce_sequential_optimization=False)`. With this setting, as many trials will be generated from the initialization stage as requested, and the optimization will move to the BayesOpt stage whenever enough trials are completed." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "[**`MaxParallelismReachedException`**](https://ax.dev/api/modelbridge.html#ax.modelbridge.generation_strategy.MaxParallelismReachedException): generation strategy restricts the number of trials that can be ran simultaneously (to encourage sequential optimization), and the parallelism limit has been reached. The correct way to handle this exception is the same as `DataRequiredError` – to wait until more trial evluations complete and log their data via `ax_client.complete_trial(...)`.\n", + " \n", + "In some cases higher parallelism is important, so `enforce_sequential_optimization=False` kwarg to AxClient allows to suppress limiting of parallelism. It's also possible to override the default parallelism setting for all stages of the optimization by passing `choose_generation_strategy_kwargs` to `ax_client.create_experiment`:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "ax_client = AxClient()\n", + "ax_client.create_experiment(\n", + " parameters=[\n", + " {\"name\": \"x\", \"type\": \"range\", \"bounds\": [-5.0, 10.0]},\n", + " {\"name\": \"y\", \"type\": \"range\", \"bounds\": [0.0, 15.0]},\n", + " ],\n", + " # Sets max parallelism to 10 for all steps of the generation strategy.\n", + " choose_generation_strategy_kwargs={\"max_parallelism_override\": 10},\n", + ")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "ax_client.get_max_parallelism() # Max parallelism is now 10 for all stages of the optimization." + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3 (ipykernel)", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.11.5" + } + }, + "nbformat": 4, + "nbformat_minor": 2 }