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)