A common issue when working on data science projects is tracking the performance of your machine learning models as you run more and more experiments. MLFlow is an open source Python library which aids developers to keep track of various experiments associated with a variety of applications including but not limited to classical machine learning, natural language processing, recommender systems, time series analysis, etc. This article will provide an in depth tutorial on how to use MLFlow to optimize the experimentation components of your projects. The following is the structure of the article :
Table of Contents
- Machine Learning Workflow
- What is MLFlow?
- Tracking
- Projects
- Models
- Registry
- Installation
- Tutorial
- Requirements
- Prediction Script
- Updating Prediction Script w/ MLFlow
- MLFlow GUI
- Concluding Remarks
- Resources
Machine Learning Workflow
The machine learning workflow generally consists of 4 steps. First is the raw data, then data preprocessing & feature engineering, followed by model development and finally deployment. As one can imagine this lifecycle is an iterative process consisting of many experiments to be run. The focus of MLFlow is primarily on the model development, you would want the best performing model to be pushed to production. Even once the best model is chosen and deployed, it’s best practice to introduce a monitoring system to track the performance of the model to identify when retraining is necessary. This iterative process occurs for every data science project and can be difficult to keep track of / maneuver when individuals are working on multiple projects at the same time or when you’re revisiting a project at a later date. This is what makes the machine learning workflow tricky and difficult, MLFlow aids in solving a lot of the pain points mentioned above.

What is MLFlow?
MLFlow is an open source library which aids data and research scientists for optimizing their machine learning workflow. It’s comprised of 4 components, tracking, projects, models and model registry. This tool has been designed to function in conjunction with any popular open source machine learning library including Sci-Kit Learn, PyTorch, Keras, TensorFlow, etc. The use of this library aids data driven researchers and scientists in making their projects reproducible and reusable by others.
Tracking
The tracking component of MLFlow refers to logging of parameters, performance / other metrics, versioning and other artifacts when building your machine learning pipeline [1]. When conducting experiments in masses, tracking of the parameters and the performance metrics associated with each experiment in an efficient manner is crucial in identifying the best performing parameters. This component is easy to set up and is compatible with both notebook and scripts. The tracking functionality of MLFlow is the bread and butter which aids in connecting the other components of MLFLow.
Projects
This component of MLFlow aims to aid its users in the reproducibility associated with the projects they’re working on. Using this component in conjunction with version control softwares like GitHub maximizes its capabilities. Using this component in conjunction with MLFlow Tracking, will allow MLFLow to remember the project versions and all parameters tracked [1]. This component holds its results in a folder usually called mlruns
which is created by MLFlow. This folder can be pushed to your GitHub repository as it shouldn’t hold any sensitive information.
Models
The modelling component to MLFlow refers to the packaging of developed models and aiding the deployment associated with them. It’s compatible with various cloud based tools including Azure, AWS (SageMaker), as well as services like Docker and Airflow. It allows models to be saved and loaded as .sav
or .pkl
files to be referenced for use later.
Registry
The registry component of MLFlow aims to aid the management of the MLFlow life cycle.
It provides model lineage (which MLflow experiment and run produced the model), model versioning, stage transitions (for example from staging to production or archiving), and annotations.
Installation
The installation process for MLFlow is quite simple, you can either rely on pip and run pip install mlflow
in your terminal / command line, or go to their GitHub repository and clone it to your computer locally. Upon installation, you can run the command mlflow --version
in terminal / command line to confirm that it has been installed locally.

Tutorial
I will outline how to update an existing model training process using MLFlow.
Requirements
The current version of Python I’m using is 3.8.8 and the following packages are required to follow along in this tutorial :
mlflow=1.25.1
numpy=1.20.1
sklearn=0.24.1
Prediction Script
The code above shows building a gradient boosting classifier for predicting the age range of users. This script will automatically generate a random dataset, train the model and evaluate the performance of the model.
Updating Prediction Script w/ MLFlow
The script above shows the updates to the prediction script for incorporating MLFlow functionality. There were a few different functions used:
-
mlflow.start_run
This function indicates the start of a new run, it will track metrics and parameters. If you want to reference an existing run then you are able to pass therun_id
associated with that run. -
mlflow.log_metric
This function will log a metric under the current active run. In the script above we log the accuracy metrics generated after training the model. mlflow.log_param
This function will log a parameter under the current active run. In the script above we log the parameters used to train the model.-
mlflow.sklearn.log_model
Log an scikit-learn model under the current run. In the script above we save the GradientBoostingClassifier model as a sklearn model. Note: There are a lot of other popular machine learning libraries supported by MLFlow. These libraries include TensorFlow, PyTorch, etc.
There are various other functions built into MLFlow not mentioned above. These are the main metrics used for any usual classification / regression related pipeline. MLFlow does provide quite extensive support for a variety of areas of Artificial Intelligence including recommendation systems, computer vision, natural language processing, etc.
Running the updated script will create a new folder in the directory called mlruns
, this folder will hold the logged values, models, etc. for each run. This allows you to compare which runs had better results given some specific parameters.

You can either manually inspect the folder or look at the results through the GUI which has a great UI.
MLFlow GUI
To initialize the GUI you simply run the command mlflow ui
in your terminal / command line. What should appear is a link to a local host server which holds the logged models, parameters and metrics associated with each run.

Going to the link http://127.0.0.1:5000
will lead you to the following page :

Clicking on any of the runs it will lead you to the following page :

You’re able to compare the parameters and metrics across all runs, and since we’ve logged the model, we’re also given quite a bit of information on how to load and use the model for future pipelines. The model is a drop-down menu consisting of model info, a yaml file associated with the conda environment, a downloadable pickle file for the model and the requirements.

This might be one of the most useful features associated with this module. After comparing the results to see which model has the best accuracy, we can easily download the model associated with the best performance. This model can then easily be fed into the production pipeline associated with the project.
Concluding Remarks
This tutorial shows how to update an existing model training and predicting pipeline using MLFlow to track various parameters, metrics and log models. It aids in the Machine Learning lifecycle associated with the user because documenting experiment(s) result(s) is a lot more efficient and easier.
Feel free to checkout the repository associated with this tutorial on my GitHub page here.
If you’re looking to transition into the data industry and want mentorship and guidance from seasoned mentors then you might want to check out Sharpest Minds. Sharpest Minds is a mentorship platform where mentors (who are seasoned practicing data scientists, machine learning engineers, research scientists, CTO, etc.) would aid in your development and learning to land a job in data. Check them out here.
Resources
If you found reading this article useful, here are others I’ve written which you might also find useful.
Comprehensive Guide to GitHub for Data Scientists
Recommendation Systems Explained
Active Learning in Machine Learning Explained
Text Summarization in Python with Jaro-Winkler and PageRank
Text Similarity w/ Levenshtein Distance in Python