Expected `choose_generation_strategy_kwargs` `num_trials` to control `optimization_complete` bool during `get_next_trials`
See original GitHub issueNot a blocker for me, but just leaving here for provenance. In practice, I’d typically defer to using AxSearch via RayTune’s interface to do asynchronous optimization, but I wanted to have an example that does batch optimization with trials in a batch running in parallel.
import ray
from ax.service.ax_client import AxClient
from ax.utils.measurement.synthetic_functions import branin
batch_size = 2
ax_client = AxClient()
ax_client.create_experiment(
parameters=[
{"name": "x1", "type": "range", "bounds": [-5.0, 10.0]},
{"name": "x2", "type": "range", "bounds": [0.0, 15.0]},
],
objective_name="branin",
minimize=True,
# Sets max parallelism to 10 for all steps of the generation strategy.
choose_generation_strategy_kwargs={
"num_trials": 21,
"max_parallelism_override": batch_size,
},
)
@ray.remote
def evaluate(parameters):
return {"branin": branin(parameters["x1"], parameters["x2"])}
optimization_complete = False
while not optimization_complete:
trial_mapping, optimization_complete = ax_client.get_next_trials(batch_size)
# start running trials in a queue (new trials will start as resources are freed)
futures = [evaluate.remote(parameters) for parameters in trial_mapping.values()]
# wait for all trials in the batch to complete before continuing (i.e. blocking)
results = ray.get(futures)
# report the completion of trials to the Ax client
for trial_index, raw_data in zip(trial_mapping.keys(), results):
ax_client.complete_trial(trial_index=trial_index, raw_data=raw_data)
I figured one workaround would be:
...
choose_generation_strategy_kwargs={
"max_parallelism_override": batch_size,
"enforce_sequential_optimization": False,
},
...
...
num_batches = np.ceil(total_trials / batch_size).astype(int)
for batch_num in range(num_batches):
current_batch_size = (
batch_size
if batch_num < num_batches - 1
else total_trials - batch_num * batch_size
)
trial_mapping, optimization_complete = ax_client.get_next_trials(current_batch_size)
...
but in the batch immediately after the transition from Sobol to GPEI, it only produces one trial instead of two.
So one real workaround is:
import numpy as np
from ax.utils.measurement.synthetic_functions import branin
batch_size = 2
total_trials = 11
ax_client = AxClient()
ax_client.create_experiment(
parameters=[
{"name": "x1", "type": "range", "bounds": [-5.0, 10.0]},
{"name": "x2", "type": "range", "bounds": [0.0, 15.0]},
],
objective_name="branin",
minimize=True,
# Sets max parallelism to 10 for all steps of the generation strategy.
choose_generation_strategy_kwargs={
"max_parallelism_override": batch_size,
"enforce_sequential_optimization": False,
},
)
@ray.remote
def evaluate(parameters):
return {"branin": branin(parameters["x1"], parameters["x2"])}
iter_num = 0
num_batches = np.ceil(total_trials / batch_size).astype(int)
for batch_num in range(num_batches):
current_batch_size = (
batch_size
if batch_num < num_batches - 1
else total_trials - batch_num * batch_size
)
trials = [ax_client.get_next_trial() for _ in range(current_batch_size)]
parameters_batch, trial_indices = zip(*trials)
# start running trials in a queue (new trials will start as resources are freed)
futures = [evaluate.remote(parameters) for parameters in parameters_batch]
# wait for all trials in the batch to complete before continuing (i.e. blocking)
results = ray.get(futures)
# report the completion of trials to the Ax client
for trial_index, raw_data in zip(trial_indices, results):
ax_client.complete_trial(trial_index=trial_index, raw_data=raw_data)
I’m guessing specifying a generation strategy directly would address this, though I don’t know if there’s still that behavior of skipping a trial during the transition from Sobol to GPEI.
Issue Analytics
- State:
- Created 10 months ago
- Comments:5 (5 by maintainers)
Top Results From Across the Web
No results found
Top Related Medium Post
No results found
Top Related StackOverflow Question
No results found
Troubleshoot Live Code
Lightrun enables developers to add logs, metrics and snapshots to live code - no restarts or redeploys required.
Start FreeTop Related Reddit Thread
No results found
Top Related Hackernoon Post
No results found
Top Related Tweet
No results found
Top Related Dev.to Post
No results found
Top Related Hashnode Post
No results found
Top GitHub Comments
Great point, we should add that to the docstring! I think the conditions are: user actually set
num_trials
on the last step of their generation strategy to something other than -1 (I think this is only possible manually, not throughchoose_generation_strategy
at the moment), search space was exhausted (applies to search spaces without range parameters), or optimization has completely converged and the model resuggested the same point more than some N times.I need to make some changes to that function soon, so can augment docstring as part of that work! Or feel free to do that yourself, a PR will be much appreciated.
Ok, I think I’m seeing my misunderstanding. There are 5 Sobol trials, so the behavior described makes sense to me. I figured it was always
2 * num_parameters
, but it seems like it’smax(2 * num_parameters, 5)
.