How can I save loss and accuracy metrics in mlflow after each epoch?

Question

I would like to see metrics like loss and accuracy as a graph by storing each value for the corresponding metrics after each epoch during training/testing phase of a keras model.

PS: I know that we can do it by using autolog feature of mlflow for keras like below, but I dont want to use that.

mlflow.keras.autolog()

Manas · Accepted Answer · 2022-07-25T12:36:34.463

After searching through the internet and combining a few concepts, I was able to solve the problem that I had asked. In Keras, we can create custom callbacks that can be called at various points (start/end of epoch, batch, etc) during training, testing, and prediction phase of a model.

So, I created a Keras custom callback to store loss/accuracy values after each epoch as mlflow metrics like below.

class CustomCallback(keras.callbacks.Callback):
    def on_epoch_end(self, epoch, logs=None):
        mlflow.log_metrics({
           "loss": logs["loss"],
           "sparse_categorical_accuracy":
               logs["sparse_categorical_accuracy"],
           "val_loss": logs["val_loss"],
           "val_sparse_categorical_accuracy":
               logs["val_sparse_categorical_accuracy"],
    })

I called this above callback during training of my model like below.

history = model.fit(
    features_train,
    labels_train,
    batch_size=BATCH_SIZE,
    epochs=EPOCHS,
    callbacks=[CustomCallback()],
    validation_split=0.2
)

The keras custom callback stored all the values during training after each epoch which I was able to see as a graph in mlflow UI like below.

How can I save loss and accuracy metrics in mlflow after each epoch?

1 Answers1