Customize metric visualization in MLFlow UI when using mlflow.tensorflow.autolog()

user_6563277 0
Viewed 1035

I'm trying to integrate MLFlow to my project. Because I'm using tf.keras.fit_generator() for my training so I take advantage of mlflow.tensorflow.autolog()(docs here) to enable automatic logging of metrics and parameters:

    model = Unet()
    optimizer = tf.keras.optimizers.Adam(LEARNING_RATE)

    metrics = [IOUScore(threshold=0.5), FScore(threshold=0.5)]
    model.compile(optimizer, customized_loss, metrics)

    callbacks = [
        tf.keras.callbacks.ModelCheckpoint("model.h5", save_weights_only=True, save_best_only=True, mode='min'),
        tf.keras.callbacks.TensorBoard(log_dir='./logs', profile_batch=0, update_freq='batch'),
    ]


    train_dataset = Dataset(src_dir=SOURCE_DIR)

    train_data_loader = DataLoader(train_dataset, BATCH_SIZE, shuffle=True)

   
    with mlflow.start_run():
        mlflow.tensorflow.autolog()
        mlflow.log_param("batch_size", BATCH_SIZE)

        model.fit_generator(
            train_data_loader,
            steps_per_epoch=len(train_data_loader),
            epochs=EPOCHS,
            callbacks=callbacks   
            )

I expected something like this (just a demonstration taken from the docs):

Visualization on the docs

However, after the training finished, this is what I got:

f1_score visualization

How can I configure so that the metric plot will update and display its value at each epoch instead of just showing the latest value?

mlflow tf.keras
1 Answers

After searching around, I found this issue related to my problem above. Actually, all my metrics just logged once each training (instead of each epoch as my intuitive thought). The reason is I didn't specify the every_n_iter parameter in mlflow.tensorflow.autolog(), which indicates how many 'iterations' must pass before MLflow logs metric executed (see the docs). So, changing my code to:

mlflow.tensorflow.autolog(every_n_iter=1)

fixed the problem.

P/s: Remember that in TF 2.x, an 'iteration' is an epoch (in TF 1.x it's a batch).

user_65632770
Related