Mlflow and KerasTuner integration

Question

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?

Answer 1

The following line of code with mlflow.start_run(run_name="myrun", nested=True) as run: is producing that each training is stored in the same "experiment".

Avoid using it and mlfow will create automatically one different experiment for each training performed by the tuner.serach