def get_untransformed_trace(experiment: Experiment, model: ModelBridge,
trace: Any):
observation_features = []
observation_data = []
metric_names = list(m for m in model.model.metric_names
if m in experiment.optimization_config.metrics)
for i, x in enumerate(trace['x']):
observation_features.append(
ObservationFeatures(parameters={
p: x[i].item()
for i, p in enumerate(model.parameters)
}))
output_means = trace['output']['mean'][i, :]
output_variances = trace['output']['variance'][i, :]
observation_data.append(
ObservationData(metric_names=metric_names,
means=output_means.numpy(),
covariance=np.diag(output_variances.numpy())))
for transform_name, transform in reversed(model.transforms.items()):
observation_features = transform.untransform_observation_features(
observation_features)
observation_data = transform.untransform_observation_data(
observation_data, observation_features)
return list(
Observation(f, d)
for f, d in zip(observation_features, observation_data))
def get_next_trial(self, model_gen_options=None):
"""
@param model_gen_options (dict or None)
Example value: {
"acquisition_function_kwargs": {
"random_scalarization": True,
},
}
"""
trial = self.experiment.new_trial(
generator_run=self.generation_strategy.gen(
experiment=self.experiment,
n=1,
pending_observations={
metric_name: [
ObservationFeatures(parameters=trial.arm.parameters,
trial_index=np.int64(trial_index))
for trial_index, trial in
self.experiment.trials.items()
if not trial.status.is_completed
]
for metric_name in self.experiment.metrics
},
model_gen_options=model_gen_options,
))
if self.verbose:
print(f"Generated Ax trial {trial.index} with model "
f"{type(self.generation_strategy.model.model).__name__}")
print(f"Marking Ax trial {trial.index} as running")
trial.mark_running(no_runner_required=True)
return trial.arm.parameters, trial.index
def init_search(self):
if self.num_sobol_trials > 0:
# BEGIN: from /ax/service/utils/dispatch.py
generation_strategy = GenerationStrategy(
name="Sobol+GPEI",
steps=[
GenerationStep(
model=Models.SOBOL,
num_trials=self.num_sobol_trials,
min_trials_observed=ceil(self.num_sobol_trials / 2),
enforce_num_trials=True,
model_kwargs={'seed': 0}
),
GenerationStep(
model=Models.GPEI,
num_trials=-1,
max_parallelism=3,
model_gen_kwargs={
"fixed_features": ObservationFeatures(
parameters={
kv["name"]: kv["value"]
for kv in self.config.get(
"ax_search.fixed_parameters"
)
}
)
},
),
],
)
# END: from /ax/service/utils/dispatch.py
self.ax_client = AxClient(generation_strategy=generation_strategy)
else:
self.ax_client = AxClient()
self.ax_client.create_experiment(
name=self.job_id,
parameters=self.config.get("ax_search.parameters"),
objective_name="metric_value",
minimize=False,
parameter_constraints=self.config.get("ax_search.parameter_constraints"),
choose_generation_strategy_kwargs={'random_seed': 0},
)
self.config.log(
"ax search initialized with {}".format(self.ax_client.generation_strategy)
)
# Make sure sobol models are resumed correctly
if self.ax_client.generation_strategy._curr.model == Models.SOBOL:
self.ax_client.generation_strategy._set_current_model(
experiment=self.ax_client.experiment, data=None
)
# Regenerate and drop SOBOL arms already generated. Since we fixed the seed,
# we will skip exactly the arms already generated in the job being resumed.
num_generated = len(self.parameters)
if num_generated > 0:
num_sobol_generated = min(
self.ax_client.generation_strategy._curr.num_arms, num_generated
)
for i in range(num_sobol_generated):
generator_run = self.ax_client.generation_strategy.gen(
experiment=self.ax_client.experiment
)
# self.config.log("Skipped parameters: {}".format(generator_run.arms))
self.config.log(
"Skipped {} of {} Sobol trials due to prior data.".format(
num_sobol_generated,
self.ax_client.generation_strategy._curr.num_arms,
)
)