mlflow.keras_core

class mlflow.keras_core.MLflowCallback(run, log_every_epoch=True, log_every_n_steps=None)

Bases: keras_core.src.callbacks.callback.Callback

Callback for logging Keras Core training metrics to MLflow.

This callback logs model metadata at training begins, and logs training metrics every epoch or every n steps (defined by the user) to MLflow.

Parameters

• run – an ‘mlflow.entities.run.Run’ instance, the MLflow run.

• log_every_epoch – bool, If True, log metrics every epoch. If False, log metrics every n steps.

• log_every_n_steps – int, log metrics every n steps. If None, log metrics every epoch. Must be None if log_every_epoch=True.

Example

import keras_core as keras
import mlflow
import numpy as np

# Prepare data for a 2-class classification.
data = np.random.uniform([8, 28, 28, 3])
label = np.random.randint(2, size=8)
model = keras.Sequential(
    [
        keras.Input([28, 28, 3]),
        keras.layers.Flatten(),
        keras.layers.Dense(2),
    ]
)
model.compile(
    loss=keras.losses.SparseCategoricalCrossentropy(from_logits=True),
    optimizer=keras.optimizers.Adam(0.001),
    metrics=[keras.metrics.SparseCategoricalAccuracy()],
)
with mlflow.start_run() as run:
    model.fit(
        data,
        label,
        batch_size=4,
        epochs=2,
        callbacks=[mlflow.keras_core.MLflowCallback(run)],
    )

on_batch_end(batch, logs=None)

Log metrics at the end of each batch with user specified frequency.

on_epoch_end(epoch, logs=None)

Log metrics at the end of each epoch.

on_test_end(logs=None)

Log validation metrics at validation end.

on_train_begin(logs=None)

Log model architecture and optimizer configuration when training begins.