Mlflow and KerasTuner integration

user_996719 0
Viewed 390

I am trying to integrate together KerasTuner and Mlflow. I'd like to record the loss at each epoch of each trial of Keras Tuner.

My approach is:

class MlflowCallback(tf.keras.callbacks.Callback):
    
    # This function will be called after each epoch.
    def on_epoch_end(self, epoch, logs=None):
        if not logs:
            return
        # Log the metrics from Keras to MLflow     
        mlflow.log_metric("loss", logs["loss"], step=epoch)
    

from kerastuner.tuners import RandomSearch

with mlflow.start_run(run_name="myrun", nested=True) as run:
  
  tuner = RandomSearch(
      train_fn,
      objective='loss',
      max_trials=25, 
  )
  tuner.search(train,
              validation_data=validation, 
              validation_steps=validation_steps,
              steps_per_epoch=steps_per_epoch, 
              epochs=5, 
              callbacks=[MlflowCallback()]
  )

However, the loss values are reported (sequentially) in one single experiment. Is there a way to record them independently? Loss values

keras-tuner mlflow keras tensorflow callback
1 Answers

The line

with `mlflow.start_run(run_name="myrun", nested=True)` as run:

is what's causing that each training is stored in the same experiment. Don't use it and mlflow will automatically create one different experiment for each training performed by the tuner.search

user_181008370
Related