Getting started with mlflow Experiment tracking

Using a local folder and SQLite

5 min readMay 30, 2022

Introduction

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

One of the features offered by mlflow is the tracking of experiments in an organized way. This post explains how to get started with this. We will consider a simple Machine Learning example and show how to track the experiments in two different ways:

Using a local folder
Using a SQLite database

There are more possibilities for tracking experiments, a detailed list can be found in the documentation of mlflow. A jupyter notebook with the example code shown here can be found on GitHub.

Getting started with mlflow

In order to use mflow, we need to install it using pip install mlflow.

Now, let us define a simple training function for a Random Forest Regressor:

def train_rf(X_train, y_train, X_val, y_val, n_estimators=100, max_depth=6):
    
    model = RandomForestRegressor(n_estimators=n_estimators, max_depth=max_depth)
    model.fit(X_train, y_train)
    
    # generate predictions
    y_pred_train = model.predict(X_train).reshape(-1,1)
    y_pred = model.predict(X_val).reshape(-1,1)
    
    # calculate errors
    rmse_train = mean_squared_error(y_pred_train, y_train, squared=False)
    rmse_val = mean_squared_error(y_pred, y_val, squared=False)
    print(f"rmse training: {rmse_train:.3f}\t rmse validation: {rmse_val:.3f}")

We can then log parameters, metrics, and much more using mlflow. This is useful to analyze how the training and evaluation metrics change and to compare different parameter setups. We will modify this function and log the parameters “n_estimators”, “max_depth” and the metric “rmse” for training and validation.

def train_rf(X_train, y_train, X_val, y_val, n_estimators=100, max_depth=6):
    
    model = RandomForestRegressor(n_estimators=n_estimators, max_depth=max_depth)
    model.fit(X_train, y_train)
    
    # generate predictions
    y_pred_train = model.predict(X_train).reshape(-1,1)
    y_pred = model.predict(X_val).reshape(-1,1)
    
    # calculate errors
    rmse_train = mean_squared_error(y_pred_train, y_train, squared=False)
    rmse_val = mean_squared_error(y_pred, y_val, squared=False)
    print(f"rmse training: {rmse_train:.3f}\t rmse validation: {rmse_val:.3f}")
    
    # Logging params and metrics to MLFlow
    mlflow.log_param('n_estimators', n_estimators)
    mlflow.log_param('max_depth', max_depth)
    mlflow.log_metric('rmse_val', rmse_val)
    mlflow.log_metric('rmse_train', rmse_train)

Track Experiments in a local Folder

The easiest way to track our experiments is to save all the logged information in a local folder. We can do this by using with mlflow.start_run() and run our model in the with statement. A directory mlflow, where the runs and all associated logs are saved is created. Every time we run code with the command with mlflow.start_run() a new run will be created under the same experiment name. By default, the experiment is named 0.

with mlflow.start_run():
    train_rf(X_train, y_train, X_val, y_val, n_estimators=100, max_depth=6)

We perform a new run of the model changing the “max_depth” parameter. If we run it as above it will be stored under the current experiment folder as the previous, in this case with experiment_id=0.

with mlflow.start_run():
    train_rf(X_train, y_train, X_val, y_val, n_estimators=100, max_depth=None)

If we want to save our results under a different experiment, we can set a new experiment, using mlflow.set_experiment(). If the experiment already exists the runs will be saved under this experiment, if it doesn't exist a new one will be created. Let‘s assume we have a second model (e.g. XGBoost) and we want to log the results under a new experiment. We need to give the experiment name, the experiment id will automatically increment.

# defining a new experiment
experiment_name = 'xgboost'
exp_id = mlflow.set_experiment(experiment_name=experiment_name)
with mlflow.start_run():
    train_xgb(X_train, y_train, X_val, y_val, learning_rate=0.1)

We can run models under specific experiment names, by setting the experiment_id to an existing experiment.

with mlflow.start_run(experiment_id=1):
    train_xgb(X_train, y_train, X_val, y_val, learning_rate=0.01)

Using tree mlruns we can see the structure of the folder, where our experiments are stored.

Structure of the folder where the experiments are stored

The numeric numbers 0, 1, … are the different experiments. In each of these folders, the runs we performed are stored — each run in a separate folder named by the experiment id. In this folder, we find all the metrics, parameters, and other things we stored.

Track Experiments using Sqlite

As an alternative, the logged results can be stored in a database. Mlflow supports MySQL, Postgres, and Sqlite. In this post, we will only consider SQlite an overview of the database tracking possibilities that can be found in the docs of mlflow. In order to tell mlflow to store things in a database, we need to set the tracking uri as mlflow.set_tracking_uri("sqlite:///mlruns.db"). Alternatively, we can export the environment variable export MLFLOW_TRACKING_URI sqlite:///mlruns.db. Then start the run as before.

mlflow.set_tracking_uri("sqlite:///mlruns.db")
with mlflow.start_run(experiment_id=0):
    train_rf(X_train, y_train, X_val, y_val)

This creates a database “mlruns.db” in the current folder. We can check the tables using SQLite, and for example, selecting the metrics stored.

Use the web UI for visualization

We can visualize the results using the web UI of mlflow to analyze and compare different runs and experiments. In the first case, when we store the logged parameters and metrics locally in .\mlflow, we can simply type mlflow ui in the terminal to see the results interactively. This works if we are at the same location where the folder mlflow was created. More generally we can reach the folder from other locations using mlflow ui --backend-store-uri file:////absolute-path/mlruns.

In the second case, we need to change the backend-store-uri and use mlflow ui --backend-store-uri sqlite:///mlruns.db. If we want to access a database from a different location we can set the path mlflow ui --backend-store-uri sqlite:////absolute-path/mlruns.db. In order for this to work sqlalchemy needs to be installed, e.g. pip install sqlalchemy.

Another useful parameter is --port to change the port for the UI, by default it is set to5000. Here is an example of what the web UI looks like. There are a lot more visualization features, this is only the starting page. Go ahead and explore them!

Get an email whenever Pumaline publishes.

Get an email whenever Pumaline publishes. By signing up, you will create a Medium account if you don't already have…

medium.com

Mlearning.ai Submission Suggestions

How to become a writer on Mlearning.ai