Supercharging Your Machine Learning Experiments with MLflow Tracking

ByJordan Brown

Machine learning is eating the world. The field has seen explosive growth over the past decade, with breakthroughs in deep learning and compute infrastructure driving a 10x increase in the number of ML papers published per year.

As ML scales in complexity and impact, the tools and processes that worked for small research projects break down. Experiment tracking becomes critical to manage the iterative process of model development. A 2021 survey by Kaggle found that 84% of data scientists consider experiment tracking to be critical or very important to their work.

While you can track experiments manually in a spreadsheet, this becomes tedious and error-prone, especially as the number of experiments grows. Fortunately, many dedicated tools now exist to streamline the process. One of the most popular is MLflow, an open source platform incubated at Databricks.

What is MLflow?

MLflow is a platform for the complete machine learning lifecycle. Since launching in 2018, it has seen rapid adoption, with over 2 million monthly downloads and contributions from over 200 companies.

The core value proposition of MLflow is to make it easy to track experiments, reproduce results, and share models across teams. It achieves this through four main components:

  1. MLflow Tracking: APIs and UI for logging parameters, metrics, artifacts, and code
  2. MLflow Projects: Package ML code in a reusable, reproducible format
  3. MLflow Models: Deploy models in diverse serving environments
  4. MLflow Model Registry: Collaborate on models and track their lineage

At its foundation is a simple yet powerful idea: every training run should be tracked in a standard format. Matei Zaharia, creator of MLflow, explains:

"Machine learning is a complex process that involves code, data, models, and hyperparameters. Getting all of these components working together requires a lot of trial and error. We created MLflow to provide a simple, standard format to track all the moving pieces and make ML experiments reproducible and collaborative."

Let‘s dive into the tracking component to see how it works in practice.

MLflow Tracking: Your Experiment Tracking Command Center

MLflow Tracking is organized around the concept of runs and experiments. An experiment is a collection of related runs, typically aimed at solving a particular problem, like customer churn prediction or object detection. Each run represents a single attempt to train a model, with a specific code version, dataset, parameters, etc.

For each run, you can log:

  • Parameters: Model hyperparameters and configuration (e.g. learning rate, batch size)
  • Metrics: Evaluation measures (e.g. accuracy, RMSE) both during and after training
  • Artifacts: Output files like trained models, plots, or data files
  • Source: The code version used to run the experiment

Logging this information is easy with MLflow‘s APIs for Python, R, Java, and REST. Here‘s a full example in Python:

import mlflow
from sklearn.neural_network import MLPClassifier

with mlflow.start_run():

    # Log parameters
    mlflow.log_param("hidden_layers", [10, 10])  
    mlflow.log_param("solver", "adam")

    # Create and train model
    model = MLPClassifier(hidden_layer_sizes=[10, 10], solver="adam")
    model.fit(X_train, y_train)

    # Log metrics over time
    for epoch in range(100):
        loss = model.loss_curve_[epoch]
        mlflow.log_metric("loss", loss, step=epoch)

    # Evaluate final metrics
    accuracy = model.score(X_test, y_test)
    mlflow.log_metric("accuracy", accuracy)

    # Log model and other artifacts 
    mlflow.sklearn.log_model(model, "model") 
    mlflow.log_artifact("features.txt")

This logs a run with the full training history, final metrics, model file, and a list of features. You can compare it with other runs in the MLflow UI:

The UI makes it easy to identify the best model and trace back to the exact code and parameters used to generate it. For more advanced workflows, you can log custom artifacts like datasets or model checkpoints to S3 or GCS buckets.

Autologging: Experiment Tracking on Autopilot

While the MLflow tracking API is straightforward, adding tracking calls throughout your code can still be tedious. That‘s where autologging comes in.

Autologging allows you to track experiments with almost no code changes. When enabled, MLflow will automatically capture the parameters, metrics, and models from popular libraries like scikit-learn, PyTorch, Keras, and XGBoost. You can enable it with a single line:

import mlflow.sklearn
mlflow.sklearn.autolog()

# Existing model code 
model = RandomForestRegressor(n_estimators=100, max_depth=5)
model.fit(X_train, y_train)

That‘s it! MLflow will now log the model type, parameters, training loss, evaluation scores, feature importance, and trained model for every run. This can save hours of manually instrumenting training scripts.

Of course, autologging has its limits. It can only track what the library exposes. For full control and customization, you‘ll want to use the manual APIs. But for many standard workflows, autologging is a huge time saver.

Tracking Large Language Models in MLflow

The past few years have seen an explosion of interest in large language models (LLMs). LLMs are massive neural networks, with billions of parameters, that can perform tasks like generation, translation, and question answering with remarkable fluency.

However, LLMs pose some unique challenges for experiment tracking:

  • Prompts and generated text are variable length and more complex than scalar metrics
  • LLMs are evaluated along many qualitative dimensions like coherence, factuality, and bias
  • Training LLMs is computationally intensive and often done in a distributed fashion

To address these challenges, MLflow 2.0 introduced a suite of features designed for tracking LLMs:

  • The mlflow.log_text API for logging prompts and generated text
  • Autologging for popular LLM libraries like HuggingFace Transformers
  • Built-in metrics like BLEU, ROUGE, and perplexity
  • Support for multi-node, multi-GPU training with tools like DeepSpeed

Here‘s an example of tracking an LLM experiment with MLflow:

import mlflow
from transformers import AutoModelForCausalLM, AutoTokenizer

model_name = "gpt2-medium"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name)

with mlflow.start_run():

    mlflow.log_param("model", model_name)
    mlflow.log_param("max_length", 30)

    prompt = "The answer to the ultimate question of life, the universe, and everything is"
    input_ids = tokenizer(prompt, return_tensors="pt").input_ids

    output = model.generate(input_ids, max_length=30, num_return_sequences=5)
    generated_text = tokenizer.batch_decode(output, skip_special_tokens=True)

    mlflow.log_text(prompt, "prompt.txt")
    mlflow.log_text("\n".join(generated_text), "generated.txt")

    perplexity = model.eval_lm(input_ids, return_dict=True)["perplexity"].item()
    mlflow.log_metric("perplexity", perplexity)

This logs an experiment with the LLM name, prompt, generated text, and perplexity. You can use this to compare different models, prompts, and decoding configurations.

The mlflow.log_text API is particularly useful for logging conversations with chatbots or question-answering systems. By tracking the full interaction history, you can debug issues and measure quality improvements over time.

With support for distributed training and specialized metrics, MLflow provides a robust foundation for tracking the next generation of language AI systems.

Collaboration and Reproducibility with MLflow

Tracking experiments is great, but it‘s even more powerful when combined with collaboration and reproducibility features. That‘s where MLflow Projects and Model Registry come in.

MLflow Projects provide a standard format to package your code, models, and data dependencies in a reproducible way. A project consists of:

  • Code to run the project
  • A MLproject config file that defines project parameters and entry points
  • An environment.yml file that lists required libraries

Projects can be run anywhere through a consistent API: locally, on a remote server, or on cloud platforms like Databricks and SageMaker. This makes it easy to share work across teams.

The MLflow Model Registry allows you to collaboratively manage models across their lifecycle. Registered models are logged with MLflow Tracking and can be tagged with stage labels like Staging and Production. This provides a full lineage of a model from experimentation to deployment.

By packaging experiments as projects and tracking deployed models in the registry, data science teams can:

  • Reproduce past experiments to validate results
  • Run experiments in different environments without code changes
  • Deploy models to production with full traceability
  • Implement approval workflows and access controls around sensitive models

Integrating MLflow with Your ML Ecosystem

A key design principle of MLflow is to be open and integrate well with other tools in the ML ecosystem. MLflow supports an extensive list of backend stores for tracking and registry data, including:

  • Local filesystem (default)
  • SQLAlchemy compatible databases (MySQL, Postgres, etc)
  • Hive Metastore
  • Databricks workspaces

You can also extend MLflow with custom plugins for new artifact stores, tracking backends, and deployment targets.

MLflow provides built-in integrations with many popular ML libraries and frameworks:

  • Scikit-Learn
  • Keras
  • PyTorch
  • XGBoost
  • LightGBM
  • Statsmodels
  • TensorFlow
  • Spark MLlib

This makes it easy to add tracking to your existing ML workflows, without having to significantly refactor code.

For a scalable ML pipeline orchestration, you can use MLflow together with Apache Airflow, Kubeflow, or Prefect. These tools provide the scheduling, resource management, and dependency resolution needed to productionize ML workflows.

Many cloud platforms also provide managed MLflow services. For example, with Databricks MLflow Experiments you can spin up a tracking server with a single click, and auto-scale training jobs using on-demand clusters. Azure ML and AWS SageMaker also have built-in support for MLflow.

By leveraging these integrations, you can build an ML platform that adapts to your specific requirements while still benefiting from MLflow‘s standardized data model and UI.

Conclusion

Experiment tracking is an essential part of professional machine learning workflows. By logging experiments in a structured format, you can iterate faster, reproduce results, and collaborate across teams.

MLflow addresses these needs through an open platform with four key components:

  1. Tracking for logging parameters, code, and results in a standard format
  2. Projects for packaging code in a reproducible way and running on any platform
  3. Models for deploying models in diverse serving environments
  4. Model Registry for collaboratively managing models across their lifecycle

Since launch, MLflow has seen significant adoption across the industry, with over 3 million monthly downloads and integrations with all major ML libraries and frameworks. It has become a foundational component of the modern data science toolkit.

As the field of ML continues to evolve, with new techniques like large language models and active learning, the importance of experiment tracking will only grow. MLflow is well poised to meet these challenges through an extensible architecture and vibrant community.

While no tool can replace rigorous statistical practices and subject matter expertise, MLflow provides the technical foundation needed to scale ML workflows to production grade systems. If you haven‘t already, I encourage you to try it out. Happy tracking!

Resources and References:

How useful was this post?

Click on a star to rate it!

Average rating 0 / 5. Vote count: 0

No votes so far! Be the first to rate this post.

You May Like to Read,

Similar Posts