ML Model Tracking

With Mlflow

Outline

  • Overview of the ML Development challenges

  • How MLflow tackles these challenges

  • MLflow components

  • Demo

  • How to get started

 

Machine Learning Development is complex

Data Preparation

Training

Deployment

Raw Data

ML Lifecycle

Data Preparation

Training

Deployment

Raw Data

ML Lifecycle

Hyperparameter
 tuning
Scale
Hyperparameter
 tuning
Scale
Scale
Scale

Model

Exchange

Custom ML Platforms 

Facebook (Meta) FBLearner, Uber Michelangelo, Google TFX 

+ They standardize the dataprep/training/deploy loop

If you work with the platform, you get these

- Limited to a few algorithms and frameworks

- Tied to a company's infrastructure

Can we provide similar benefits in an open manner? 

Introducing Mlflow 

Open Machine Learning Platform

  • Works with any ML library and language

  • Runs the same way anywhere (e.g any cloud)

  • Designed to be useful for 1 or 1000+ organisations

Mlflow Components

Tracking

Projects

Models

Record and query experiments, code, configs, results...etc

Packaging format for reproducible runs on any platform

General model format that supports diverse deployment tools

Tracking

  • Parameters: Key-value pairs of your code

  • Metrics: Numeric values (can update over time)

  • Source: Training code that ran

  • Version: Version of the training code

  • Artefacts: Files, including data and models

  • Tags and Notes: Any additional information

Tracking

Tracking server

Tracking

import mlflow

with mlflow.start_run():
    mlflow.log_param("layers", layers)
    mlflow.log_param("alpha", alpha)


   # Train Model 
   mlflow.log_metric("mse", model.mse)
   mlflow.log_artefact("plot", model.plot)
   mlflow.tensorflow.log_model("model")

Tracking

Record and query experiments, code, configs, results...etc

Tracking Demo

Keras Text Classification Example

Mlflow Components

Tracking

Projects

Models

Record and query experiments, code, configs, results...etc

Packaging format for reproducible runs on any platform

General model format that supports diverse deployment tools

Projects Motivation

Diverse set of training tools

Diverse set of environments

Challenge: ML results are Hard to reproduce

Project spec

Data

Code

config

Local Execution

Remote Execution

Projects

Packaging project for reproducible ML Runs

  • Any code folder or github repo
  • Optional MLProject file with code spec

Defines dependencies for reproducibility

  • Conda (R, Docker etc) can be specified in MLProject
  • Reproducible in (almost) every environment

Execution API for running projects

  • CLI/Python/Java/R
  • Supports local and remote execution

Projects

Example

my_project/
MLProject
conda.yaml
main.py
model.py
...
conda_env: conda.yaml

entry_points:
  main:
    parameters:
      training_data: path

    command: |
          python main.py {training_data}

Projects

Example

Demo

Tracking

Projects

Models

Record and query experiments, code, configs, results...etc

Packaging format for reproducible runs on any platform

General model format that supports diverse deployment tools

Models Motivation

ML Frameworks 

Serving tools 

Inference code

Batch and stream scoring

Models Motivation

ML Frameworks 

Serving tools 

Inference code

Batch and stream scoring

Model format

Flavor 1

Flavor 2

e.g Pytorch

e.g Tensorflow

ML Frameworks 

Serving tools 

Inference code

Batch and stream scoring

Standard for ML Models

MLFlow Models 

Packaging format for ML Models

  • Any directory with MLmodel file 

Defines dependencies for reproducibility

  • Conda environment can be specified in MLmodel configuration

Model creation utilities

  • Save models from any framework in MLflow format

Deployment APIs

  • CLI / Python/ R/ Java

Example MLflow Model


 mlflow.tensorflow.log_model(...)

my_model/
MLmodel
estimator/
saved_model.pb
...
variables/

Example MLflow Model


run_id: 4785fd4c4bb663461
time_created: 2023-06-28T12:34
flavors:
	tensorflow:
    	saved_model_dir: estimator
        signature_def_key:predict
    python_function:
    	loader_module:
            mlflow.tensorflow
my_model/
MLmodel
estimator/
saved_model.pb
...
variables/

 mlflow.tensorflow.log_model(...)

}

}

usable within Tensorflow tools/APIs
usable within any Python tool

MLflow Flavors Example

Model format

Flavor1:

Pyfunc

Flavor2:

Keras

Train a model

mlflow.keras.log_model()
predict = mlflow.pyfunc.load_pyfunc()
predict(input_dataframe)

model = mlflow.keras.load_model()
model.predict(keras.Input(..))

MLflow Flavors Example

predict = mlflow.pyfunc.load_pyfunc()
predict(input_dataframe)

MLFlow Models Demo

Recap

  • Overview of the ML Development challenges

  • How MLflow tackles these challenges

  • MLflow components

  • Demo

  • How to get started

 


  pip install mlflow

Conclusion

  • Newer components; Model Registry for Model Management

  • Multi-step project workflows 

  • Fluent tracking API for scala and Java

  • Better environment support when loading schemas 

  • Improved model input/output schemas

 

Thank you