How to manage an end-to-end machine learning project with MLflow? part 1

Have you ever felt bored when working on a machine learning project and it is struggling to manage the model version? It could be much more difficult when members of a data team are working on the same project at the same time.

Which is the latest training dataset?

Which is the best-performing model?

Many questions can pop up along the way and sometimes it is hard to answer those questions. Isn’t it better to have a tool that can manage each element in the machine learning project life cycle like datasets and models, so a data team doesn’t have to worry about the tedious tasks and can focus more on value-added tasks?

After knowing MLflow, I can manage machine learning projects much easier. In this article, I will share my experience with MLflow with code in Python that you can adapt to your case.

Contents of this article

• What is MLflow?
• Why use MLflow?
• MLflow concept
• MLflow tracking example

What is MLflow?

MLflow is an open-source platform to manage the ML lifecycle, including experimentation, reproducibility, deployment, and a central model registry.
(from the MLflow official website)

There are over 13 million monthly downloads!!
MLflow has become the standard platform for end-to-end MLOps.
You can read my article about MLOps here.

MLflow main advantages including:

• It works with many machine-learning libraries such as scikit-learn, TensorFlow, and Keras.
• It works with many platforms such as Databrick, Azure, and AWS.
• It can be scaled to big data.
• It is easy to collaborate within a team (small or big).
• It is open-source.
• The MLflow developer team is very active.

Why use MLflow?

• Don’t know which ml model version is used in production.
  > MLflow can help in model management.
• Have been trying many models and features but can’t track each model's performance. It isn't easy to compare the model’s performance.
  > MLflow offers experiment tracking.
• Don’t have a system to keep model visualization images such as a confusion matrix.
  > MLflow has artifact storage.
• A code works only on my computer but not others.
  > MLflow makes code reusable and reproducible.
• I want to access the model of others to try my own experiment.
  > MLflow allows collaboration within a team.
• Model training is completed but you don’t know how to implement the model in the real world. Don’t know how to track a model’s performance after deployment.
  > MLflow makes the model deployment easier and trackable.

MLflow concept

MLflow components

MLflow consists of 4 compartments including

1. Tracking: logging parameters, code, metrics, and files such as images, and data sets.
2. Models: packaging machine learning models for later use by downstream tools.
3. Model Registry: a centralized model store for team collaboration.
4. Projects: packaging data science code in a reusable and reproducible way.

Some MLflow functions can be implemented via API or via UI on a web browser.

Let’s start coding!!

First of all, you have to set up MLflow on your computer.

It is recommended to create a new python environment for MLflow using Anaconda.

Open the Anaconda prompt and enter the below code line by line.

conda create --name mlflow
conda activate mlflow
pip install mlflow

Now you finish creating a new python environment for MLflow. You are ready to get your hand dirty.

💡You can learn more about the python environment setting in the “Install Python” part of this article.

MLflow tracking example

MLflow will keep tracking on the concept of “run”.

Each run will track the Code version, Start and end time of the run, Source, Key-value input parameters, Key-value metrics, and Artifacts.
Recorded artifacts can be in any format such as images or data files.

Several runs can be grouped into an “experiment” to easily manage the same task.

MLflow tracking steps

1. Set where your run data is recorded (tracking URI)
2. Create an experiment
3. Start run
4. Log parameters/metrics/artifacts, add tags

Let’s see where your run data is recorded.

MLflow runs can be recorded to local files, databases, or remote servers.
By default, the MLflow will create a new folder called mlruns in your working directory and log your runs there.

mlruns folder

To log runs on the other place, you have to set the MLFLOW_TRACKING_URI environment variable to a tracking server’s URI or call mlflow.set_tracking_uri().

These are what tracking URIs will look like for each option:

• Local file pathfile:/my/local/dir
• Database <dialect>+<driver>://<username>:<password>@<host>:<port>/<database>
• HTTP server https://my-server:5000
• Databricks workspace databricks or databricks://<profileName>

Code example

In this example, I will train a logistic regression model to classify the (well-known) iris dataset. Then, I will log the parameters, and metrics of this model along with the model itself using MLflow. After that, I will tell you how to load those logged elements.

• Import libraries

import os
import pandas as pd
import mlflow
from sklearn.datasets import load_iris
from sklearn.model_selection import train_test_split
from sklearn.linear_model import LogisticRegression
from sklearn.metrics import accuracy_score,precision_score,recall_score,log_loss

• Load dataset and prepare data

iris = load_iris()
X = pd.DataFrame(iris.data, columns=iris.feature_names)
y = iris.target

X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.3,stratify = y, random_state=55)

X, y for model training

• Train logistic regression model on the training data set

classifier = LogisticRegression()
classifier.fit(X_train,y_train)

• Make a prediction on the test set and evaluate the model.
  Save evaluation metrics as a dictionary. We will save this metrics dictionary for this run using mlflow.

y_pred = classifier.predict(X_test)

acc = accuracy_score(y_test, y_pred)
prec = precision_score(y_test, y_pred,average='micro')
recall = recall_score(y_test, y_pred,average='micro')
metrics = {'accuracy': round(acc, 2), 'precision': round(prec, 2), 'recall': round(recall, 2)}

metrics

• Define model parameters to be logged in the run such as coefficient, and intercept in this example.

coef = classifier.coef_
intercept = classifier.intercept_
params = {'coef':coef, 'intercept': intercept}

params

• Create a visualization image to be saved as an artifact. I created a scatter plot in this example and saved it as a png file. Actually, you can save any file type as an artifact.

import seaborn as sns
import matplotlib.pyplot as plt

sns.scatterplot(x=X['sepal length (cm)'], y=X['sepal width (cm)'], hue=y, palette="deep")
plt.savefig('scatter.png')

Scatter plot

It is time for the mlflow!!

• Setting tracking_uri, experiment name, and run name
  I set a tracking URI on my local computer (localhost), so other users will not be able to access my mlflow run data.
  You can set a tracking URI to the database as well, so your team can access the same MLflow data source.

mlflow.set_tracking_uri("http://localhost:5000")
experiment_name = "iris_classifier"
run_name="iris_classifier_ver1.0"

• Log run data
  ⚠️Before running this code you may have to call mlflow ui on the Anaconda Prompt first or you may face a server access error.

mlflow.set_experiment(experiment_name)
with mlflow.start_run(run_name=run_name):
    mlflow.log_params(params)
    mlflow.log_metrics(metrics)
    mlflow.log_artifact('./scatter.png')
    mlflow.sklearn.log_model(classifier, "model")
print('Run - %s is logged to Experiment - %s' %(run_name, experiment_name))

Code explanation

• mlflow.set_experiment: set an experiment as active. If the experiment does not exist, create a new experiment.
• mlflow.start_run: new runs are launched under this experiment.
• mlflow.log_params: Log key-value parameters in the currently active run. If you use log_param without “s”, you can log only 1 key-value parameter
• mlflow.log_metrics: Log key-value parameters in the currently active run. If you use log_param without “s”, you can log only 1 key-value parameter
• mlflow.log_artifacts: Log files in a given directory as artifacts (an image file in this example). You can set the artifact path to be different from the tracking server.
• mlflow.log_model: Log a trained model in the currently active run. I will talk about this in more detail in the mlflow model component (maybe in the part 2 article).

— End of the code —

Let’s see the result of our tracking.

MLflow UI

After tracking, you can see the logged data in a beautiful user interface by calling the below code on the Anaconda Prompt.

mlflow ui

Go to http://127.0.0.1:5000/ on your web browser which is the localhost server and you will be something like the below picture.

MLflow UI

You will see the experiment you created, iris_classifier. Click it.

Inside the experiment, You will see a list of runs. Click on the run to see more detail in that run.

You can find your logged parameters, metrics, artifacts, and model.

Run detail

You can also add tags for each run (you can learn from the MLflow official website)

There will be many files auto-generated in the model folder. I will not talk about these files in this article.

You can learn more from the official document about MLflow tracking

Load logged elements

It will not be helpful if the logged metrics, parameters, and artifacts can’t be loaded to use again.

These are codes that will be used to load logged elements.

• Get a list of all experiments.

experiments = mlflow.search_experiments()
for exp in experiments:
    print(exp.name)

result

• Get info on a specific experiment.

experiment_name = "iris_classifier"
mlflow.get_experiment_by_name(experiment_name)

result

• Get a list of runs. The result is a Pandas Dataframe, so you can manipulate it using pandas functions.

runs = mlflow.search_runs(experiment_names=[experiment_name])
runs

This is a list of columns of the runs dataframe. It depends on your metrics and parameters.

[‘run_id’, ‘experiment_id’, ‘status’, ‘artifact_uri’, ‘start_time’,
‘end_time’, ‘metrics.precision’, ‘metrics.accuracy’, ‘metrics.recall’,
‘params.coef’, ‘params.intercept’, ‘tags.mlflow.source.type’,
‘tags.mlflow.source.name’, ‘tags.mlflow.runName’, ‘tags.mlflow.user’,
‘tags.mlflow.log-model.history’]

Next, I will store the run id of the latest run for later use.

run_id = runs.iloc[0]["run_id"]

• Show artifacts list

client = mlflow.tracking.MlflowClient()
client.list_artifacts(run_id = run_id)

result

There are 2 artifacts in my run including model and scatter.png.

• Load artifacts

download_path = "C:/Users/yanna/OneDrive/Desktop" # \ is an escape character, so use / instead.
client.download_artifacts(run_id = run_id, path="scatter.png", dst_path = download_path)

I downloaded a scatter.png to my local computer. You can also load it to the online server. (Don’t forget to change the download_path to your destination folder path.)

• Load a model and make a prediction

logged_model = 'runs:/'+str(run_id)+'/model' # specify model from run id
loaded_model = mlflow.pyfunc.load_model(model_uri=logged_model)
prediction = loaded_model.predict(X_test)
print(prediction)

result

In this code, I loaded the model from the mlflow tracking server by specifying the run id. Then, I make a prediction of the X_test using the loaded model. The result shown is the predicted classes of each row of the X_test.

You can find this example code in my GitHub.

This is a basic of MLflow tracking. There are much more features of MLflow which you can explore from the official document.

Here is the list of MLflow features you might be interested in.

• Record run and artifact separately on the servers.
• Automatic logging using mlflow.autolog().
• Add tags to run.

Conclusion

In this article, you learn about MLflow which is a platform for machine learning project management. MLflow including tracking, models, model registry, and project. I cover only tracking in part 1. After finishing training the model, you can set tracking URI and create runs inside an experiment.
In a run, you can log parameters, metrics, and artifacts. Then, you can load your logged data from your run when needed.

If this article is helpful, please clap and follow me for more data science articles.

BECOME a WRITER at MLearning.ai

Mlearning.ai Submission Suggestions