def test_raw_data_format_with_fidelities(self):
ax_client = AxClient()
ax_client.create_experiment(
parameters=[
{"name": "x", "type": "range", "bounds": [-5.0, 10.0]},
{"name": "y", "type": "range", "bounds": [0.0, 1.0]},
],
minimize=True,
)
for _ in range(6):
parameterization, trial_index = ax_client.get_next_trial()
x, y = parameterization.get("x"), parameterization.get("y")
ax_client.complete_trial(
trial_index,
raw_data=[
({"y": y / 2.0}, {"objective": (branin(x, y / 2.0), 0.0)}),
({"y": y}, {"objective": (branin(x, y), 0.0)}),
],
)
def test_verify_parameterization(self):
ax_client = AxClient()
ax_client.create_experiment(
name="test_experiment",
parameters=[
{
"name": "x",
"type": "range",
"bounds": [-5.0, 10.0]
},
{
"name": "y",
"type": "range",
"bounds": [0.0, 15.0]
},
],
minimize=True,
objective_name="a",
)
params, trial_idx = ax_client.get_next_trial()
self.assertTrue(
ax_client.verify_trial_parameterization(trial_index=trial_idx,
parameterization=params))
# Make sure it still works if ordering in the parameterization is diff.
self.assertTrue(
ax_client.verify_trial_parameterization(
trial_index=trial_idx,
parameterization={
k: params[k]
for k in reversed(list(params.keys()))
},
))
self.assertFalse(
ax_client.verify_trial_parameterization(
trial_index=trial_idx,
parameterization={k: v + 1.0
for k, v in params.items()},
))
def get_one_batch_of_trials(
ax_client: AxClient,
parallelism: Tuple[int, int],
num_trials_so_far: int,
num_max_trials_to_do: int,
) -> TrialBatch:
"""Returns a TrialBatch that contains a list of trials that can be
run in parallel. TrialBatch also flags if the search space is exhausted."""
is_search_space_exhausted = False
# Ax throws an exception if the search space is exhausted. We catch
# the exception and set the flag to True
(num_trials, max_parallelism_setting) = parallelism
if max_parallelism_setting == -1:
# Special case, we can group all the trials into one batch
max_parallelism_setting = num_trials - num_trials_so_far
if num_trials == -1:
# This is a special case where we can run as many trials in parallel as we want.
# Given that num_trials is also -1, we can run all the trials in parallel.
max_parallelism_setting = num_max_trials_to_do
list_of_trials = []
for _ in range(max_parallelism_setting):
try:
parameters, trial_index = ax_client.get_next_trial()
list_of_trials.append(
Trial(
overrides=map_params_to_arg_list(params=parameters),
trial_index=trial_index,
)
)
except SearchSpaceExhausted:
is_search_space_exhausted = True
break
return TrialBatch(
list_of_trials=list_of_trials,
is_search_space_exhausted=is_search_space_exhausted,
)
def test_annotate_exception(self, _):
ax_client = AxClient()
ax_client.create_experiment(
name="test_experiment",
parameters=[
{"name": "x", "type": "range", "bounds": [-5.0, 10.0]},
{"name": "y", "type": "range", "bounds": [0.0, 15.0]},
],
minimize=True,
objective_name="a",
)
with self.assertRaisesRegex(
expected_exception=RuntimeError,
expected_regex="Cholesky errors typically occur",
):
ax_client.get_next_trial()
def test_unnamed_experiment_snapshot(self):
ax_client = AxClient(random_seed=239)
ax_client.create_experiment(parameters=[
{
"name": "x",
"type": "range",
"bounds": [-5.0, 10.0]
},
{
"name": "y",
"type": "range",
"bounds": [0.0, 15.0]
},
])
serialized = ax_client.to_json_snapshot()
ax_client = AxClient.from_json_snapshot(serialized)
self.assertIsNone(ax_client.experiment._name)
def test_start_and_end_time_in_trial_completion(self):
start_time = current_timestamp_in_millis()
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]},
],
minimize=True,
)
params, idx = ax_client.get_next_trial()
ax_client.complete_trial(
trial_index=idx,
raw_data=1.0,
metadata={
"start_time": start_time,
"end_time": current_timestamp_in_millis(),
},
)
dat = ax_client.experiment.fetch_data().df
self.assertGreater(dat["end_time"][0], dat["start_time"][0])
def test_default_generation_strategy_discrete(self) -> None:
"""Test that Sobol is used if no GenerationStrategy is provided and
the search space is discrete.
"""
# Test that Sobol is chosen when all parameters are choice.
ax_client = AxClient()
ax_client.create_experiment(
parameters=[ # pyre-fixme[6]: expected union that should include
{"name": "x", "type": "choice", "values": [1, 2, 3]},
{"name": "y", "type": "choice", "values": [1, 2, 3]},
]
)
self.assertEqual(
[s.model for s in not_none(ax_client.generation_strategy)._steps],
[Models.SOBOL],
)
self.assertEqual(ax_client.get_max_parallelism(), [(-1, -1)])
self.assertTrue(ax_client.get_trials_data_frame().empty)
def test_log_failure(self):
ax = AxClient()
ax.create_experiment(
parameters=[
{
"name": "x1",
"type": "range",
"bounds": [-5.0, 10.0]
},
{
"name": "x2",
"type": "range",
"bounds": [0.0, 15.0]
},
],
minimize=True,
)
_, idx = ax.get_next_trial()
ax.log_trial_failure(idx, metadata={"dummy": "test"})
self.assertTrue(ax.experiment.trials.get(idx).status.is_failed)
self.assertEqual(
ax.experiment.trials.get(idx).run_metadata.get("dummy"), "test")
def test_fail_on_batch(self):
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]},
],
minimize=True,
)
batch_trial = ax_client.experiment.new_batch_trial(
generator_run=GeneratorRun(
arms=[
Arm(parameters={"x1": 0, "x2": 1}),
Arm(parameters={"x1": 0, "x2": 1}),
]
)
)
with self.assertRaises(NotImplementedError):
ax_client.complete_trial(batch_trial.index, 0)
def test_storage_error_handling(self, mock_save_fails):
"""Check that if `suppress_storage_errors` is True, AxClient won't
visibly fail if encountered storage errors.
"""
init_test_engine_and_session_factory(force_init=True)
config = SQAConfig()
encoder = Encoder(config=config)
decoder = Decoder(config=config)
db_settings = DBSettings(encoder=encoder, decoder=decoder)
ax_client = AxClient(db_settings=db_settings, suppress_storage_errors=True)
ax_client.create_experiment(
name="test_experiment",
parameters=[
{"name": "x", "type": "range", "bounds": [-5.0, 10.0]},
{"name": "y", "type": "range", "bounds": [0.0, 15.0]},
],
minimize=True,
)
for _ in range(3):
parameters, trial_index = ax_client.get_next_trial()
ax_client.complete_trial(
trial_index=trial_index, raw_data=branin(*parameters.values())
)
def test_attach_trial_and_get_trial_parameters(self):
ax_client = AxClient()
ax_client.create_experiment(
parameters=[
{
"name": "x",
"type": "range",
"bounds": [-5.0, 10.0]
},
{
"name": "y",
"type": "range",
"bounds": [0.0, 15.0]
},
],
minimize=True,
)
params, idx = ax_client.attach_trial(parameters={"x": 0.0, "y": 1.0})
ax_client.complete_trial(trial_index=idx, raw_data=5)
self.assertEqual(ax_client.get_best_parameters()[0], params)
self.assertEqual(ax_client.get_trial_parameters(trial_index=idx), {
"x": 0,
"y": 1
})
with self.assertRaises(ValueError):
ax_client.get_trial_parameters(
trial_index=10) # No trial #10 in experiment.
with self.assertRaisesRegex(ValueError, ".* is of type"):
ax_client.attach_trial({"x": 1, "y": 2})
def test_ttl_trial(self):
ax_client = AxClient()
ax_client.create_experiment(
parameters=[
{
"name": "x",
"type": "range",
"bounds": [-5.0, 10.0]
},
{
"name": "y",
"type": "range",
"bounds": [0.0, 15.0]
},
],
minimize=True,
)
# A ttl trial that ends adds no data.
params, idx = ax_client.get_next_trial(ttl_seconds=1)
self.assertTrue(ax_client.experiment.trials.get(idx).status.is_running)
time.sleep(1) # Wait for TTL to elapse.
self.assertTrue(ax_client.experiment.trials.get(idx).status.is_failed)
# Also make sure we can no longer complete the trial as it is failed.
with self.assertRaisesRegex(
ValueError,
".* has been marked FAILED, so it no longer expects data."):
ax_client.complete_trial(trial_index=idx,
raw_data={"objective": (0, 0.0)})
params2, idy = ax_client.get_next_trial(ttl_seconds=1)
ax_client.complete_trial(trial_index=idy, raw_data=(-1, 0.0))
self.assertEqual(ax_client.get_best_parameters()[0], params2)
def test_abandon_trial(self):
ax_client = AxClient()
ax_client.create_experiment(
parameters=[
{
"name": "x",
"type": "range",
"bounds": [-5.0, 10.0]
},
{
"name": "y",
"type": "range",
"bounds": [0.0, 15.0]
},
],
minimize=True,
)
# An abandoned trial adds no data.
params, idx = ax_client.get_next_trial()
ax_client.abandon_trial(trial_index=idx)
data = ax_client.experiment.fetch_data()
self.assertEqual(len(data.df.index), 0)
# Can't update a completed trial.
params2, idx2 = ax_client.get_next_trial()
ax_client.complete_trial(trial_index=idx2,
raw_data={"objective": (0, 0.0)})
with self.assertRaisesRegex(ValueError, ".* in a terminal state."):
ax_client.abandon_trial(trial_index=idx2)
def test_trial_completion(self):
ax_client = AxClient()
ax_client.create_experiment(
parameters=[
{
"name": "x",
"type": "range",
"bounds": [-5.0, 10.0]
},
{
"name": "y",
"type": "range",
"bounds": [0.0, 15.0]
},
],
minimize=True,
)
params, idx = ax_client.get_next_trial()
# Can't update before completing.
with self.assertRaisesRegex(ValueError, ".* not yet"):
ax_client.update_trial_data(trial_index=idx,
raw_data={"objective": (0, 0.0)})
ax_client.complete_trial(trial_index=idx,
raw_data={"objective": (0, 0.0)})
# Cannot complete a trial twice, should use `update_trial_data`.
with self.assertRaisesRegex(ValueError, ".* already been completed"):
ax_client.complete_trial(trial_index=idx,
raw_data={"objective": (0, 0.0)})
# Cannot update trial data with observation for a metric it already has.
with self.assertRaisesRegex(ValueError, ".* contained an observation"):
ax_client.update_trial_data(trial_index=idx,
raw_data={"objective": (0, 0.0)})
# Same as above, except objective name should be getting inferred.
with self.assertRaisesRegex(ValueError, ".* contained an observation"):
ax_client.update_trial_data(trial_index=idx, raw_data=1.0)
ax_client.update_trial_data(trial_index=idx, raw_data={"m1": (1, 0.0)})
metrics_in_data = ax_client.experiment.fetch_data(
).df["metric_name"].values
self.assertIn("m1", metrics_in_data)
self.assertIn("objective", metrics_in_data)
self.assertEqual(ax_client.get_best_parameters()[0], params)
params2, idy = ax_client.get_next_trial()
ax_client.complete_trial(trial_index=idy, raw_data=(-1, 0.0))
self.assertEqual(ax_client.get_best_parameters()[0], params2)
params3, idx3 = ax_client.get_next_trial()
ax_client.complete_trial(trial_index=idx3,
raw_data=-2,
metadata={"dummy": "test"})
self.assertEqual(ax_client.get_best_parameters()[0], params3)
self.assertEqual(
ax_client.experiment.trials.get(2).run_metadata.get("dummy"),
"test")
best_trial_values = ax_client.get_best_parameters()[1]
self.assertEqual(best_trial_values[0], {"objective": -2.0})
self.assertTrue(
math.isnan(best_trial_values[1]["objective"]["objective"]))
def test_keep_generating_without_data(self):
# Check that normally numebr of arms to generate is enforced.
ax_client = AxClient()
ax_client.create_experiment(
parameters=[
{
"name": "x",
"type": "range",
"bounds": [-5.0, 10.0]
},
{
"name": "y",
"type": "range",
"bounds": [0.0, 15.0]
},
],
minimize=True,
)
for _ in range(5):
parameterization, trial_index = ax_client.get_next_trial()
with self.assertRaisesRegex(DataRequiredError,
"All trials for current model"):
ax_client.get_next_trial()
# Check thatwith enforce_sequential_optimization off, we can keep
# generating.
ax_client = AxClient(enforce_sequential_optimization=False)
ax_client.create_experiment(
parameters=[
{
"name": "x",
"type": "range",
"bounds": [-5.0, 10.0]
},
{
"name": "y",
"type": "range",
"bounds": [0.0, 15.0]
},
],
minimize=True,
)
self.assertFalse(
ax_client.generation_strategy._steps[0].enforce_num_trials, False)
self.assertFalse(
ax_client.generation_strategy._steps[1].max_parallelism, None)
for _ in range(10):
parameterization, trial_index = ax_client.get_next_trial()
def test_create_experiment(self) -> None:
"""Test basic experiment creation."""
ax_client = AxClient(
GenerationStrategy(
steps=[GenerationStep(model=Models.SOBOL, num_trials=30)]))
with self.assertRaisesRegex(ValueError,
"Experiment not set on Ax client"):
ax_client.experiment
ax_client.create_experiment(
name="test_experiment",
parameters=[
{
"name": "x",
"type": "range",
"bounds": [0.001, 0.1],
"value_type": "float",
"log_scale": True,
},
{
"name": "y",
"type": "choice",
"values": [1, 2, 3],
"value_type": "int",
"is_ordered": True,
},
{
"name": "x3",
"type": "fixed",
"value": 2,
"value_type": "int"
},
{
"name": "x4",
"type": "range",
"bounds": [1.0, 3.0],
"value_type": "int",
},
{
"name": "x5",
"type": "choice",
"values": ["one", "two", "three"],
"value_type": "str",
},
{
"name": "x6",
"type": "range",
"bounds": [1.0, 3.0],
"value_type": "int",
},
],
objective_name="test_objective",
minimize=True,
outcome_constraints=["some_metric >= 3", "some_metric <= 4.0"],
parameter_constraints=["x4 <= x6"],
)
assert ax_client._experiment is not None
self.assertEqual(ax_client._experiment, ax_client.experiment)
self.assertEqual(
ax_client._experiment.search_space.parameters["x"],
RangeParameter(
name="x",
parameter_type=ParameterType.FLOAT,
lower=0.001,
upper=0.1,
log_scale=True,
),
)
self.assertEqual(
ax_client._experiment.search_space.parameters["y"],
ChoiceParameter(
name="y",
parameter_type=ParameterType.INT,
values=[1, 2, 3],
is_ordered=True,
),
)
self.assertEqual(
ax_client._experiment.search_space.parameters["x3"],
FixedParameter(name="x3",
parameter_type=ParameterType.INT,
value=2),
)
self.assertEqual(
ax_client._experiment.search_space.parameters["x4"],
RangeParameter(name="x4",
parameter_type=ParameterType.INT,
lower=1.0,
upper=3.0),
)
self.assertEqual(
ax_client._experiment.search_space.parameters["x5"],
ChoiceParameter(
name="x5",
parameter_type=ParameterType.STRING,
values=["one", "two", "three"],
),
)
self.assertEqual(
ax_client._experiment.optimization_config.outcome_constraints[0],
OutcomeConstraint(
metric=Metric(name="some_metric"),
op=ComparisonOp.GEQ,
bound=3.0,
relative=False,
),
)
self.assertEqual(
ax_client._experiment.optimization_config.outcome_constraints[1],
OutcomeConstraint(
metric=Metric(name="some_metric"),
op=ComparisonOp.LEQ,
bound=4.0,
relative=False,
),
)
self.assertTrue(
ax_client._experiment.optimization_config.objective.minimize)
def test_default_generation_strategy_continuous(self, _a, _b, _c,
_d) -> None:
"""Test that Sobol+GPEI is used if no GenerationStrategy is provided."""
ax_client = AxClient()
ax_client.create_experiment(
parameters=[ # pyre-fixme[6]: expected union that should include
{
"name": "x",
"type": "range",
"bounds": [-5.0, 10.0]
},
{
"name": "y",
"type": "range",
"bounds": [0.0, 15.0]
},
],
objective_name="a",
minimize=True,
)
self.assertEqual(
[s.model for s in not_none(ax_client.generation_strategy)._steps],
[Models.SOBOL, Models.GPEI],
)
with self.assertRaisesRegex(ValueError, ".* no trials"):
ax_client.get_optimization_trace(objective_optimum=branin.fmin)
for i in range(6):
parameterization, trial_index = ax_client.get_next_trial()
x, y = parameterization.get("x"), parameterization.get("y")
ax_client.complete_trial(
trial_index,
raw_data={
"a": (
checked_cast(
float,
branin(checked_cast(float, x),
checked_cast(float, y)),
),
0.0,
)
},
sample_size=i,
)
self.assertEqual(ax_client.generation_strategy.model._model_key,
"GPEI")
ax_client.get_optimization_trace(objective_optimum=branin.fmin)
ax_client.get_contour_plot()
ax_client.get_feature_importances()
trials_df = ax_client.get_trials_data_frame()
self.assertIn("x", trials_df)
self.assertIn("y", trials_df)
self.assertIn("a", trials_df)
self.assertEqual(len(trials_df), 6)
def test_plotting_validation(self):
ax_client = AxClient()
ax_client.create_experiment(parameters=[{
"name": "x3",
"type": "fixed",
"value": 2,
"value_type": "int"
}])
with self.assertRaisesRegex(ValueError, ".* there are no trials"):
ax_client.get_contour_plot()
with self.assertRaisesRegex(ValueError, ".* there are no trials"):
ax_client.get_feature_importances()
ax_client.get_next_trial()
with self.assertRaisesRegex(ValueError, ".* less than 2 parameters"):
ax_client.get_contour_plot()
ax_client = AxClient()
ax_client.create_experiment(parameters=[
{
"name": "x",
"type": "range",
"bounds": [-5.0, 10.0]
},
{
"name": "y",
"type": "range",
"bounds": [0.0, 15.0]
},
])
ax_client.get_next_trial()
with self.assertRaisesRegex(ValueError, "If `param_x` is provided"):
ax_client.get_contour_plot(param_x="y")
with self.assertRaisesRegex(ValueError, "If `param_x` is provided"):
ax_client.get_contour_plot(param_y="y")
with self.assertRaisesRegex(ValueError, 'Parameter "x3"'):
ax_client.get_contour_plot(param_x="x3", param_y="x3")
with self.assertRaisesRegex(ValueError, 'Parameter "x4"'):
ax_client.get_contour_plot(param_x="x", param_y="x4")
with self.assertRaisesRegex(ValueError, 'Metric "nonexistent"'):
ax_client.get_contour_plot(param_x="x",
param_y="y",
metric_name="nonexistent")
with self.assertRaisesRegex(UnsupportedPlotError,
"Could not obtain contour"):
ax_client.get_contour_plot(param_x="x",
param_y="y",
metric_name="objective")
with self.assertRaisesRegex(ValueError, "Could not obtain feature"):
ax_client.get_feature_importances()
def test_overwrite(self):
init_test_engine_and_session_factory(force_init=True)
ax_client = AxClient()
ax_client.create_experiment(
name="test_experiment",
parameters=[
{
"name": "x",
"type": "range",
"bounds": [-5.0, 10.0]
},
{
"name": "y",
"type": "range",
"bounds": [0.0, 15.0]
},
],
minimize=True,
)
# Log a trial
parameters, trial_index = ax_client.get_next_trial()
ax_client.complete_trial(trial_index=trial_index,
raw_data=branin(*parameters.values()))
with self.assertRaises(ValueError):
# Overwriting existing experiment.
ax_client.create_experiment(
name="test_experiment",
parameters=[
{
"name": "x",
"type": "range",
"bounds": [-5.0, 10.0]
},
{
"name": "y",
"type": "range",
"bounds": [0.0, 15.0]
},
],
minimize=True,
)
# Overwriting existing experiment with overwrite flag.
ax_client.create_experiment(
name="test_experiment",
parameters=[
{
"name": "x1",
"type": "range",
"bounds": [-5.0, 10.0]
},
{
"name": "x2",
"type": "range",
"bounds": [0.0, 15.0]
},
],
overwrite_existing_experiment=True,
)
# There should be no trials, as we just put in a fresh experiment.
self.assertEqual(len(ax_client.experiment.trials), 0)
# Log a trial
parameters, trial_index = ax_client.get_next_trial()
self.assertIn("x1", parameters.keys())
self.assertIn("x2", parameters.keys())
ax_client.complete_trial(trial_index=trial_index,
raw_data=branin(*parameters.values()))
def test_attach_trial_ttl_seconds(self):
ax_client = AxClient()
ax_client.create_experiment(
parameters=[
{
"name": "x",
"type": "range",
"bounds": [-5.0, 10.0]
},
{
"name": "y",
"type": "range",
"bounds": [0.0, 15.0]
},
],
minimize=True,
)
params, idx = ax_client.attach_trial(parameters={
"x": 0.0,
"y": 1.0
},
ttl_seconds=1)
self.assertTrue(ax_client.experiment.trials.get(idx).status.is_running)
time.sleep(1) # Wait for TTL to elapse.
self.assertTrue(ax_client.experiment.trials.get(idx).status.is_failed)
# Also make sure we can no longer complete the trial as it is failed.
with self.assertRaisesRegex(
ValueError,
".* has been marked FAILED, so it no longer expects data."):
ax_client.complete_trial(trial_index=idx, raw_data=5)
params2, idx2 = ax_client.attach_trial(parameters={
"x": 0.0,
"y": 1.0
},
ttl_seconds=1)
ax_client.complete_trial(trial_index=idx2, raw_data=5)
self.assertEqual(ax_client.get_best_parameters()[0], params2)
self.assertEqual(ax_client.get_trial_parameters(trial_index=idx2), {
"x": 0,
"y": 1
})
def test_trial_completion(self):
ax = AxClient()
ax.create_experiment(
parameters=[
{
"name": "x1",
"type": "range",
"bounds": [-5.0, 10.0]
},
{
"name": "x2",
"type": "range",
"bounds": [0.0, 15.0]
},
],
minimize=True,
)
params, idx = ax.get_next_trial()
ax.complete_trial(trial_index=idx, raw_data={"objective": (0, 0.0)})
self.assertEqual(ax.get_best_parameters()[0], params)
params2, idx2 = ax.get_next_trial()
ax.complete_trial(trial_index=idx2, raw_data=(-1, 0.0))
self.assertEqual(ax.get_best_parameters()[0], params2)
params3, idx3 = ax.get_next_trial()
ax.complete_trial(trial_index=idx3,
raw_data=-2,
metadata={"dummy": "test"})
self.assertEqual(ax.get_best_parameters()[0], params3)
self.assertEqual(
ax.experiment.trials.get(2).run_metadata.get("dummy"), "test")
self.assertEqual(
ax.get_best_parameters()[1],
({
"objective": -2.0
}, {
"objective": {
"objective": 0.0
}
}),
)
def test_keep_generating_without_data(self):
# Check that normally numebr of arms to generate is enforced.
ax = AxClient()
ax.create_experiment(
parameters=[
{
"name": "x1",
"type": "range",
"bounds": [-5.0, 10.0]
},
{
"name": "x2",
"type": "range",
"bounds": [0.0, 15.0]
},
],
minimize=True,
)
for _ in range(5):
parameterization, trial_index = ax.get_next_trial()
with self.assertRaisesRegex(ValueError,
"All trials for current model"):
ax.get_next_trial()
# Check thatwith enforce_sequential_optimization off, we can keep
# generating.
ax = AxClient(enforce_sequential_optimization=False)
ax.create_experiment(
parameters=[
{
"name": "x1",
"type": "range",
"bounds": [-5.0, 10.0]
},
{
"name": "x2",
"type": "range",
"bounds": [0.0, 15.0]
},
],
minimize=True,
)
for _ in range(10):
parameterization, trial_index = ax.get_next_trial()
def ml_run(self, run_id=None):
seed_randomness(self.random_seed)
mlflow.log_params(flatten(get_params_of_task(self)))
total_training_time = 0
# should land to 'optimizer_props'
params_space = [
{
'name': 'lr',
'type': 'range',
'bounds': [1e-6, 0.008],
# 'value_type': 'float',
'log_scale': True,
},
{
'name': 'beta_1',
'type': 'range',
'bounds': [.0, 0.9999],
'value_type': 'float',
# 'log_scale': True,
},
{
'name': 'beta_2',
'type': 'range',
'bounds': [.0, 0.9999],
'value_type': 'float',
# 'log_scale': True,
}
]
# TODO: make reproducibility of search
# without it we will get each time new params
# for example we can use:
# ax.storage.sqa_store.structs.DBSettings
# DBSettings(url="sqlite://<path-to-file>")
# to store experiments
ax = AxClient(
# can't use that feature yet.
# got error
# NotImplementedError:
# Saving and loading experiment in `AxClient` functionality currently under development.
# db_settings=DBSettings(url=self.output()['ax_settings'].path)
)
# FIXME: temporal solution while ax doesn't have api to (re-)store state
class_name = get_class_name_as_snake(self)
ax.create_experiment(
name=f'{class_name}_experiment',
parameters=params_space,
objective_name='score',
minimize=should_minimize(self.metric),
# parameter_constraints=['x1 + x2 <= 2.0'], # Optional.
# outcome_constraints=['l2norm <= 1.25'], # Optional.
)
trial_index = 0
experiment = self._get_ax_experiment()
if experiment:
print('AX: restore experiment')
print('AX: num_trials:', experiment.num_trials)
ax._experiment = experiment
trial_index = experiment.num_trials - 1
model_task = get_model_task_by_name(self.model_name)
while trial_index < self.max_runs:
print(f'AX: Running trial {trial_index + 1}/{self.max_runs}...')
# get last unfinished trial
parameters = get_last_unfinished_params(ax)
if parameters is None:
print('AX: generate new Trial')
parameters, trial_index = ax.get_next_trial()
# good time to store experiment (with new Trial)
with self.output()['ax_experiment'].open('w') as f:
print('AX: store experiment: ', ax.experiment)
pickle.dump(ax.experiment, f)
print('AX: parameters', parameters)
# now is time to evaluate model
model_result = yield model_task(
parent_run_id=run_id,
random_seed=self.random_seed,
# TODO: actually we should be able to pass even nested params
# **parameters,
optimizer_props=parameters)
# TODO: store run_id in Trial
model_run_id = self.get_run_id_from_result(model_result)
with model_result['metrics'].open('r') as f:
model_metrics = yaml.load(f)
model_score_mean = model_metrics[self.metric]['val']
# TODO: we might know it :/
model_score_error = 0.0
total_training_time += model_metrics['train_time']['total']
with model_result['params'].open('r') as f:
model_params = yaml.load(f)
print('AX: complete trial:', trial_index)
ax.complete_trial(
trial_index=trial_index,
raw_data={'score': (model_score_mean, model_score_error)},
metadata={
'metrics': model_metrics,
'params': model_params,
'run_id': model_run_id,
})
best_parameters, _ = ax.get_best_parameters()
mlflow.log_metric('train_time.total', total_training_time)
print('best params', best_parameters)
best_trial = get_best_trial(experiment, self.metric)
mlflow.log_metrics(flatten(best_trial.run_metadata['metrics']))
mlflow.log_params(flatten(best_trial.run_metadata['params']))
class IterativePrune:
def __init__(self):
self.parser_args = None
self.ax_client = None
self.base_model_path = "base_model"
self.pruning_amount = None
def run_mnist_model(self, base=False):
parser_dict = vars(self.parser_args)
if base:
mlflow.start_run(run_name="BaseModel")
mlflow.pytorch.autolog()
dm = MNISTDataModule(**parser_dict)
dm.setup(stage="fit")
model = LightningMNISTClassifier(**parser_dict)
trainer = pl.Trainer.from_argparse_args(self.parser_args)
trainer.fit(model, dm)
trainer.test()
if os.path.exists(self.base_model_path):
shutil.rmtree(self.base_model_path)
mlflow.pytorch.save_model(trainer.get_model(), self.base_model_path)
return trainer
def load_base_model(self):
path = Path(_download_artifact_from_uri(self.base_model_path))
model_file_path = os.path.join(path, "data/model.pth")
return torch.load(model_file_path)
def initialize_ax_client(self):
self.ax_client = AxClient()
self.ax_client.create_experiment(
parameters=[{
"name": "amount",
"type": "range",
"bounds": [0.05, 0.15],
"value_type": "float"
}],
objective_name="test_accuracy",
)
@staticmethod
def prune_and_save_model(model, amount):
for _, module in model.named_modules():
if isinstance(module, torch.nn.Conv2d) or isinstance(
module, torch.nn.Linear):
prune.l1_unstructured(module, name="weight", amount=amount)
prune.remove(module, "weight")
mlflow.pytorch.save_state_dict(model.state_dict(), ".")
model = torch.load("state_dict.pth")
os.remove("state_dict.pth")
return model
@staticmethod
def count_model_parameters(model):
table = PrettyTable(["Modules", "Parameters"])
total_params = 0
for name, parameter in model.named_parameters():
if not parameter.requires_grad:
continue
param = parameter.nonzero(as_tuple=False).size(0)
table.add_row([name, param])
total_params += param
return table, total_params
@staticmethod
def write_prune_summary(summary, params):
tempdir = tempfile.mkdtemp()
try:
summary_file = os.path.join(tempdir, "pruned_model_summary.txt")
params = "Total Trainable Parameters :" + str(params)
with open(summary_file, "w") as f:
f.write(str(summary))
f.write("\n")
f.write(str(params))
try_mlflow_log(mlflow.log_artifact, local_path=summary_file)
finally:
shutil.rmtree(tempdir)
def iterative_prune(self, model, parametrization):
if not self.pruning_amount:
self.pruning_amount = parametrization.get("amount")
else:
self.pruning_amount += 0.15
mlflow.log_metric("PRUNING PERCENTAGE", self.pruning_amount)
pruned_model = self.prune_and_save_model(model, self.pruning_amount)
model.load_state_dict(copy.deepcopy(pruned_model))
summary, params = self.count_model_parameters(model)
self.write_prune_summary(summary, params)
trainer = self.run_mnist_model()
metrics = trainer.callback_metrics
test_accuracy = metrics.get("avg_test_acc")
return test_accuracy
def initiate_pruning_process(self, model):
total_trials = int(vars(self.parser_args)["total_trials"])
trial_index = None
for i in range(total_trials):
parameters, trial_index = self.ax_client.get_next_trial()
print(
"***************************************************************************"
)
print("Running Trial {}".format(i + 1))
print(
"***************************************************************************"
)
with mlflow.start_run(nested=True, run_name="Iteration" + str(i)):
mlflow.set_tags({"AX_TRIAL": i})
# calling the model
test_accuracy = self.iterative_prune(model, parameters)
# completion of trial
self.ax_client.complete_trial(trial_index=trial_index,
raw_data=test_accuracy.item())
# Ending the Base run
mlflow.end_run()
def get_parser_args(self):
parser = argparse.ArgumentParser()
parser = pl.Trainer.add_argparse_args(parent_parser=parser)
parser = LightningMNISTClassifier.add_model_specific_args(
parent_parser=parser)
parser.add_argument(
"--total_trials",
default=3,
help=
"Number of AX trials to be run for the optimization experiment",
)
self.parser_args = parser.parse_args()
def test_recommended_parallelism(self):
ax_client = AxClient()
with self.assertRaisesRegex(ValueError, "No generation strategy"):
ax_client.get_max_parallelism()
ax_client.create_experiment(
parameters=[
{
"name": "x",
"type": "range",
"bounds": [-5.0, 10.0]
},
{
"name": "y",
"type": "range",
"bounds": [0.0, 15.0]
},
],
minimize=True,
)
self.assertEqual(ax_client.get_max_parallelism(), [(5, 5), (-1, 3)])
self.assertEqual(
run_trials_using_recommended_parallelism(
ax_client, ax_client.get_max_parallelism(), 20),
0,
)
# With incorrect parallelism setting, the 'need more data' error should
# still be raised.
ax_client = AxClient()
ax_client.create_experiment(
parameters=[
{
"name": "x",
"type": "range",
"bounds": [-5.0, 10.0]
},
{
"name": "y",
"type": "range",
"bounds": [0.0, 15.0]
},
],
minimize=True,
)
with self.assertRaisesRegex(DataRequiredError,
"All trials for current model "):
run_trials_using_recommended_parallelism(ax_client, [(6, 6),
(-1, 3)], 20)
def test_default_generation_strategy(self) -> None:
"""Test that Sobol+GPEI is used if no GenerationStrategy is provided."""
ax = AxClient()
ax.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,
)
self.assertEqual(
[s.model for s in ax.generation_strategy._steps],
[Models.SOBOL, Models.GPEI],
)
for _ in range(6):
parameterization, trial_index = ax.get_next_trial()
x1, x2 = parameterization.get("x1"), parameterization.get("x2")
ax.complete_trial(trial_index,
raw_data={"branin": (branin(x1, x2), 0.0)})
# Test that Sobol is chosen when all parameters are choice.
ax = AxClient()
ax.create_experiment(parameters=[
{
"name": "x1",
"type": "choice",
"values": [1, 2, 3]
},
{
"name": "x2",
"type": "choice",
"values": [1, 2, 3]
},
])
self.assertEqual([s.model for s in ax.generation_strategy._steps],
[Models.SOBOL])
self.assertEqual(ax.get_recommended_max_parallelism(), [(-1, -1)])
def test_sqa_storage(self):
init_test_engine_and_session_factory(force_init=True)
config = SQAConfig()
encoder = Encoder(config=config)
decoder = Decoder(config=config)
db_settings = DBSettings(encoder=encoder, decoder=decoder)
ax_client = AxClient(db_settings=db_settings)
ax_client.create_experiment(
name="test_experiment",
parameters=[
{
"name": "x",
"type": "range",
"bounds": [-5.0, 10.0]
},
{
"name": "y",
"type": "range",
"bounds": [0.0, 15.0]
},
],
minimize=True,
)
for _ in range(5):
parameters, trial_index = ax_client.get_next_trial()
ax_client.complete_trial(trial_index=trial_index,
raw_data=branin(*parameters.values()))
gs = ax_client.generation_strategy
ax_client = AxClient(db_settings=db_settings)
ax_client.load_experiment_from_database("test_experiment")
# Trial #4 was completed after the last time the generation strategy
# generated candidates, so pre-save generation strategy was not
# "aware" of completion of trial #4. Post-restoration generation
# strategy is aware of it, however, since it gets restored with most
# up-to-date experiment data. Do adding trial #4 to the seen completed
# trials of pre-storage GS to check their equality otherwise.
gs._seen_trial_indices_by_status[TrialStatus.COMPLETED].add(4)
self.assertEqual(gs, ax_client.generation_strategy)
with self.assertRaises(ValueError):
# Overwriting existing experiment.
ax_client.create_experiment(
name="test_experiment",
parameters=[
{
"name": "x",
"type": "range",
"bounds": [-5.0, 10.0]
},
{
"name": "y",
"type": "range",
"bounds": [0.0, 15.0]
},
],
minimize=True,
)
with self.assertRaises(ValueError):
# Overwriting existing experiment with overwrite flag with present
# DB settings. This should fail as we no longer allow overwriting
# experiments stored in the DB.
ax_client.create_experiment(
name="test_experiment",
parameters=[{
"name": "x",
"type": "range",
"bounds": [-5.0, 10.0]
}],
overwrite_existing_experiment=True,
)
# Original experiment should still be in DB and not have been overwritten.
self.assertEqual(len(ax_client.experiment.trials), 5)
def testConvertAx(self):
from ray.tune.suggest.ax import AxSearch
from ax.service.ax_client import AxClient
config = {
"a": tune.sample.Categorical([2, 3, 4]).uniform(),
"b": {
"x": tune.sample.Integer(0, 5).quantized(2),
"y": 4,
"z": tune.sample.Float(1e-4, 1e-2).loguniform()
}
}
converted_config = AxSearch.convert_search_space(config)
ax_config = [
{
"name": "a",
"type": "choice",
"values": [2, 3, 4]
},
{
"name": "b/x",
"type": "range",
"bounds": [0, 5],
"value_type": "int"
},
{
"name": "b/y",
"type": "fixed",
"value": 4
},
{
"name": "b/z",
"type": "range",
"bounds": [1e-4, 1e-2],
"value_type": "float",
"log_scale": True
},
]
client1 = AxClient(random_seed=1234)
client1.create_experiment(parameters=converted_config)
searcher1 = AxSearch(ax_client=client1)
client2 = AxClient(random_seed=1234)
client2.create_experiment(parameters=ax_config)
searcher2 = AxSearch(ax_client=client2)
config1 = searcher1.suggest("0")
config2 = searcher2.suggest("0")
self.assertEqual(config1, config2)
self.assertIn(config1["a"], [2, 3, 4])
self.assertIn(config1["b"]["x"], list(range(5)))
self.assertEqual(config1["b"]["y"], 4)
self.assertLess(1e-4, config1["b"]["z"])
self.assertLess(config1["b"]["z"], 1e-2)
searcher = AxSearch(metric="a", mode="max")
analysis = tune.run(_mock_objective,
config=config,
search_alg=searcher,
num_samples=1)
trial = analysis.trials[0]
assert trial.config["a"] in [2, 3, 4]
mixed_config = {"a": tune.uniform(5, 6), "b": tune.uniform(8, 9)}
searcher = AxSearch(space=mixed_config, metric="a", mode="max")
config = searcher.suggest("0")
self.assertTrue(5 <= config["a"] <= 6)
self.assertTrue(8 <= config["b"] <= 9)
"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],
},
]
client = AxClient(enforce_sequential_optimization=False)
client.create_experiment(
parameters=parameters,
objective_name="hartmann6",
minimize=True, # Optional, defaults to False.
parameter_constraints=["x1 + x2 <= 2.0"], # Optional.
outcome_constraints=["l2norm <= 1.25"], # Optional.
)
algo = AxSearch(client, max_concurrent=4)
scheduler = AsyncHyperBandScheduler(reward_attr="hartmann6")
run(easy_objective, name="ax", search_alg=algo, **config)
def test_fixed_random_seed_reproducibility(self):
ax_client = AxClient(random_seed=239)
ax_client.create_experiment(parameters=[
{
"name": "x",
"type": "range",
"bounds": [-5.0, 10.0]
},
{
"name": "y",
"type": "range",
"bounds": [0.0, 15.0]
},
])
for _ in range(5):
params, idx = ax_client.get_next_trial()
ax_client.complete_trial(idx,
branin(params.get("x"), params.get("y")))
trial_parameters_1 = [
t.arm.parameters for t in ax_client.experiment.trials.values()
]
ax_client = AxClient(random_seed=239)
ax_client.create_experiment(parameters=[
{
"name": "x",
"type": "range",
"bounds": [-5.0, 10.0]
},
{
"name": "y",
"type": "range",
"bounds": [0.0, 15.0]
},
])
for _ in range(5):
params, idx = ax_client.get_next_trial()
ax_client.complete_trial(idx,
branin(params.get("x"), params.get("y")))
trial_parameters_2 = [
t.arm.parameters for t in ax_client.experiment.trials.values()
]
self.assertEqual(trial_parameters_1, trial_parameters_2)
