MLflow Example#

MLflow is a platform to streamline machine learning development, including tracking experiments, packaging code into reproducible runs, and sharing and deploying models.

Flyte provides an easy-to-use interface to log the task’s metrics and parameters to either Flyte Deck or MLflow server.

Let’s first import the libraries.

from flytekit import task, workflow
from flytekitplugins.mlflow import mlflow_autolog
import mlflow.keras
import tensorflow as tf

Run a model training here and generate metrics and parameters. Add mlflow_autolog to the task, then flyte will automatically log the metric to the Flyte Deck.

def train_model(epochs: int):
    # Refer to
    fashion_mnist = tf.keras.datasets.fashion_mnist
    (train_images, train_labels), (_, _) = fashion_mnist.load_data()
    train_images = train_images / 255.0

    model = tf.keras.Sequential([
        tf.keras.layers.Flatten(input_shape=(28, 28)),
        tf.keras.layers.Dense(128, activation='relu'),

                  metrics=['accuracy']), train_labels, epochs=epochs)
Model Metrics
Model Parameters

Finally, we put everything together into a workflow:

def ml_pipeline(epochs: int):

if __name__ == "__main__":
    print(f"Running {__file__} main...")

Total running time of the script: ( 0 minutes 0.000 seconds)

Gallery generated by Sphinx-Gallery