Multi-task BO with Service API #2546
Comments
|
Great question @sgbaird! This is something we do internally in the service API with batch trials. We do have future plans to open source our With batch trials, you'd have a GS consisting of
What are you trying to do? I noticed honegumi in the prompt. Is this for the honegumi interface or a real world use case? Is it intentional or accidental that this use case has non batch trials? |
|
Hi @sgbaird. Multi-task BO can take many forms depending on what you're trying to achieve. In any case, you need to provide Ax with a way to identify what task each trial belongs to. One way to do this would be to add a task parameter in your search space. When generating new trials, you can specify what task you want them to be generated as well. Here's an example using import numpy as np
from ax.core.observation import ObservationFeatures
from ax.modelbridge.generation_strategy import GenerationStep, GenerationStrategy
from ax.modelbridge.registry import Models
from ax.modelbridge.transforms.task_encode import TaskChoiceToIntTaskChoice
from ax.modelbridge.transforms.unit_x import UnitX
from ax.service.ax_client import AxClient, ObjectiveProperties
from ax.utils.common.typeutils import not_none
from ax.utils.measurement.synthetic_functions import hartmann6
from ax.utils.notebook.plotting import init_notebook_plotting, render
init_notebook_plotting()
# Update as needed. See ax/modelbridge/registry for default list of transforms.
transforms= [
TaskChoiceToIntTaskChoice, # Since we're using a string valued task parameter.
UnitX,
]
# Custom generation strategy to support the multi-task search space.
generation_strategy = GenerationStrategy(
name="MultiTaskMBM",
steps=[
GenerationStep(
model=Models.SOBOL,
num_trials=5,
model_kwargs={"deduplicate": True, "transforms": transforms},
),
GenerationStep(
model=Models.BOTORCH_MODULAR,
num_trials=-1,
model_kwargs={"transforms": transforms},
),
],
)
ax_client = AxClient(generation_strategy=generation_strategy)
ax_client.create_experiment(
name="hartmann_test_experiment",
parameters=[
{
"name": "x1",
"type": "range",
"bounds": [0.0, 1.0],
},
{
"name": "x2",
"type": "range",
"bounds": [0.0, 1.0],
},
{
"name": "x3",
"type": "range",
"bounds": [0.0, 1.0],
},
{
"name": "x4",
"type": "range",
"bounds": [0.0, 1.0],
},
{
"name": "x5",
"type": "range",
"bounds": [0.0, 1.0],
},
{
"name": "x6",
"type": "range",
"bounds": [0.0, 1.0],
},
# Add the task parameter!
{
"name": "task",
"type": "choice",
"values": ["base", "shifted"],
"is_task": True,
"target_value": "base"
}
],
objectives={"hartmann6": ObjectiveProperties(minimize=True)},
)
# Evaluation produces different results based on task value.
def evaluate(parameterization):
x = np.array([parameterization.get(f"x{i+1}") for i in range(6)])
value = hartmann6(x)
if parameterization.get("task") == "shifted":
value += 100
# In our case, standard error is 0, since we are computing a synthetic function.
return {"hartmann6": (value, 0.0)}
for i in range(10):
trial = ax_client.experiment.new_trial(
generator_run=ax_client.generation_strategy.gen(
experiment=ax_client.experiment,
n=1,
pending_observations=ax_client._get_pending_observation_features(
experiment=ax_client.experiment
),
# Need to specify what task we want to generate from. Switching between the two here.
fixed_features=ObservationFeatures(
{"task": "base" if i % 2 else "shifted"}
),
)
)
trial.mark_running(no_runner_required=True)
parameterization, trial_index = not_none(trial.arm).parameters, trial.index
ax_client.complete_trial(
trial_index=trial_index, raw_data=evaluate(parameterization)
)
# We can verify that the model is a ModelListGP of MultiTaskGP.
mb = ax_client.generation_strategy.model
mb.model.surrogate.model |
Sorry this took so long to get back to you! I lost track of this. Yes, the idea was for Honegumi and a BO tutorial on multi-task, since there are a lot of chemistry and materials science use-cases like this. I was trying to keep it in simplest case, so I hadn't considered/included batch trials. If there's not an easy way to support batch trials, then I'd set it so the |
Thank you for this! I was able to run it and plan to do some additional testing. |
|
@sgbaird is this still open or resolved? : ) |
|
Will consider solved for now, and post back/reopen if I run into any issues! Thanks all for the help! 🙂 |
Should I perish the thought?
xref: #1038
Fairly naive, non-functional starting point:
Would this be something for
BOTORCH_MODULARin a custom generation strategy instead?The text was updated successfully, but these errors were encountered: