The world’s leading publication for data science, AI, and ML professionals.

Using MLFlow to Track and Version Machine Learning Models

A step-by-step coding practice

When we tune parameters of a machine learning model, we might need to train it for multiple times in order to pick the best model. If the number of trainings goes to too many, we might have two problems.

  • How do we track the parameters and metrics of each model? Writing them manually to an excel file? That would be tedious and error-prone.
  • How do we version each model? Save them to the disk with a different name for each? It would be hard to remember which model came from what parameters.

MLFlow is what we need to solve these problems. It is a powerful MLOps tool for ML model tracking, versioning, packaging, and so on. In this blog, we will focus on tracking and versioning. Regarding tracking, MLFlow can track parameters, metrics, and models. Regarding versioning, MLFlow stores models in a model registry, and then users can easily choose a specific version.

MLFlow can be run locally or from a docker container, and also can be deployed to kubernetes. It has APIs in Python, Java, R, and REST.

In this blog, we will showcase how to use mlflow to track and version mnist classification models. We will first run an MNIST example with tensorflow, and then extend the code to integrate mlflow.

Image by Isaac Smith from Unsplash
Image by Isaac Smith from Unsplash

MNIST example

First, create a virtual environment, and pip install tensorflow and mlflow.

# my version is 2.6.0
pip install tensorflow
# my version is 4.4.0
pip install tensorflow_datasets
# my version is 1.9.1
pip install mlflow

Below is the code to run mnist classification using tensorflow.

Run the code, and you should see a log like below. It also saves the models in the model folder.

Log of mnist.py (Image by author)
Log of mnist.py (Image by author)

MLFlow tracking

Now, let’s track the parameters, metrics, and also artifacts (models). See the code below, (also see complete code at the bottom). First, we need to name a run. We can even name an experiment (higher level of runs) if we want. Then, we use the functions log_param , log_metric , and log_artifacts to log parameters, metrics, and artifacts.

import mlflow
with mlflow.start_run(run_name=run_name):
  mlflow.log_param("batch_size", batch_size)
  mlflow.log_param("learning_rate", learning_rate)
  mlflow.log_param("epochs", epochs)
  mlflow.log_metric("train_loss", train_loss)
  mlflow.log_metric("train_accuracy", train_acc)
  mlflow.log_metric("val_loss", val_loss)
  mlflow.log_metric("val_accuracy", val_acc)
  mlflow.log_artifacts("./model")

After running the mlflow code, we can see there is a new folder called mlruns in our disk. This is the place where parameters, metrics, and artifacts are stored locally.

Then, we can do some parameters tuning, for example, change the batch_size and learning_rate. Each run will be logged to mlflow.

Now, let’s view all runs at mlflow’s UI. Within the virtual environment, type:

mlflow ui

Then, browse at http://localhost:5000/. We can see all runs are logged there. The parameters, metrics, and run names can be clearly seen in one page.

Runs on mlflow tracking (Image by author)
Runs on mlflow tracking (Image by author)

If we click into a run, we can see more details about this run. Apart from parameters and metrics, the section of Artifacts at bottom-left shows our artifacts (models).

Details of a mlflow run (Image by author)
Details of a mlflow run (Image by author)

The url address of each run has the following format. Run ids within each experiment are unique.

http://localhost:5000/#/experiments/<experiment id>/runs/<run id>

MLFlow model versioning

MLflow Model Registry is a central location to store and version models. With it, a model has an iterative version from (for example) v1, v2, …, to v10. For each model and version, we can write a markdown description (for example detailed parameters) along with it, so that we know later what the model represents. In addition, we can tag a version with either Staging , Production , or Archived .

To set up model registry, we need a database backend in order to store models, see here for instructions. After that, we can upload our tensorflow models to the mlflow registry. Here below is a code example.

import mlflow.tensorflow
from tensorflow.python.saved_model import signature_constants
tag=[tf.saved_model.tag_constants.SERVING]
key=signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY
mlflow.tensorflow.log_model(tf_saved_model_dir="./model",
                            tf_meta_graph_tags=tag,
                            tf_signature_def_key=key,
                            artifact_path="model",
                            registered_model_name="mnist")

In conclusion, MLFlow is a powerful MLOps tool to track and version Machine Learning models. It supports APIs in python, java, R, etc. With its nice UI, we can clearly track, store, and compare different model versions. MLFlow has become a popular MLOps tool in many ML projects.

Below is the complete code for mnist classification in tensorflow with mlflow tracking.


Related Articles