[Question] Migrate hyperparameter optimization from SimpleExperiment to Experiment about ax

[lena-kashtelyan]

Hi @surajgattani, can you just use the Ax Service API (AxClient, tutorial: https://ax.dev/tutorials/gpei_hartmann_service.html)? That should resolve all the issues you are facing.

[surajgattani]

Thanks a lot @lena-kashtelyan -- Appreciate the quick response
I will try this out and post any questions that I might have!

[lena-kashtelyan]

How can I update the save and load functions to maintain backwards compatibility with the old json files?

Do you mean you'd like to be able to load old experiments essentially? You should be able to load them as AxClient.load_from_json_file. Let me know how that works out for you!

[surajgattani]

@lena-kashtelyan I was able to solve all the issues with the save and load experiment from json store.
Also, the example in the tutorial -- link explains the custom metric definition as well which allowed me to replace the evaluation function from the previous implementation.
I changed the custom definition a tiny bit marked with --> which worked for me:

from ax import Data
import pandas as pd
from ax.utils.common.result import Ok

class BoothMetric(Metric):
    def fetch_trial_data(self, trial):
        records = []
        for arm_name, arm in trial.arms_by_name.items():
            params = arm.parameters
            records.append(
                {
                    "arm_name": arm_name,
                    "metric_name": self.name,
                    "trial_index": trial.index,
                    # in practice, the mean and sem will be looked up based on trial metadata
                    # but for this tutorial we will calculate them
                    "mean": (params["x1"] + 2 * params["x2"] - 7) ** 2
                    + (2 * params["x1"] + params["x2"] - 5) ** 2,
                    "sem": 0.0,
                }
            )
        ---> return Ok(value=Data(df=pd.DataFrame.from_records(records)))

    def is_available_while_running(self) -> bool:
        return True

Thanks a lot for your quick response! Appreciate it