keras-tuner and MLFlow integration

[pbischoff]

Is there going to be an option to track hyperparameter tuning with keras-tuner by mlflow? Not sure how that would be implemented, but I guess by either enabling custom objective functions or implementing a callback which handles mlflow.

[jonnor]

The KerasTuner API has an integration point designed for things like that. One can pass a Logger instance when setting up the search. Unfortunately it does not seem to be documented much. But you can see the interface here: https://github.com/keras-team/keras-tuner/blob/master/keras_tuner/engine/logger.py#L16

Create a sublass of that, and call mlflow from there. Example for logging the selected hyperparameters:

    def register_trial(self, trial_id, trial_state):
        hyperparams = trial_state['hyperparameters']['values']
        mlflow.log_params(hyperparams)

[jonnor]

One limitation is that there does not seem to be a hook for each execution.

Adding calls to the logger for start and end execution to the self.executions_per_trial loop would be good. https://github.com/keras-team/keras-tuner/blob/3952464e4c904c942ebdc46145105aa625ed2605/keras_tuner/engine/tuner.py#L285 The end call should probably receive the history object. Both calls should preferably include the execution number as an id.

[fzyzcjy]

hi is there any updates? thanks!

[jonnor]

Here is some code I made to integrate mlflow and keras-tuner: https://gist.github.com/jonnor/c107f3ca24a36c91d8ff94029a0cd357

I think that the code in LoggerTunerMixin should be added into keras-tuner (improvements to logging hooks). And the MlflowLogger into mlflow.

[abdulbasitds]

@jonnor, what does this do? does this display hyperparameter choices in UI?

[jonnor]

@abdulbasitds yes, it tracks parameters and metrics. It uses nested mlflow runs to keep track of each kearas tuner search, each trial in the search, and each execution of each trial. Unfortunately, the UI does not show the nesting so clearly, but it is in the data structure at least. For now, there tracking of metrics and parameters is on the execution (each training run) - it is not merged up to the trial or the search yet.

See this screenshot for an example:

[abdulbasitds]

@abdulbasitds yes, it tracks parameters and metrics. It uses nested mlflow runs to keep track of each kearas tuner search, each trial in the search, and each execution of each trial. Unfortunately, the UI does not show the nesting so clearly, but it is in the data structure at least. For now, there tracking of metrics and parameters is on the execution (each training run) - it is not merged up to the trial or the search yet.

See this screenshot for an example:

Thanks @jonnor , great

[rakshithvsk]

Hi @jonnor, This is great!! It is in fact a good work around to capture metrics per trial.

However I noticed a major issue. If my epochs per trial is >1, ideally it should be persisting the best epoch out of all the epochs within a trial. But looks like its persisting the latest epoch per trial, though it not the best one. By any chance did you come across the same issue? Any suggestions would help. Thanks!!

[jonnor]

Hi @rakshithvsk, that could very well be. I do not have time right now to have a look at it. If you find a fix, please do post the updated code as Gist :)

[NusretOzates]

Hey everyone, maybe this approach could help.

I wrote the article while trying to do same thing