Quick start Python

MLflow is an open source platform for managing the end-to-end machine learning lifecycle. MLflow provides simple APIs for logging metrics (for example, model loss), parameters (for example, learning rate), and fitted models, making it easy to analyze training results or deploy models later on.

This article provides a Python quickstart, explaining how to:

  • (1 min) Install MLflow
  • (1 min) Use MLflow’s one-line autolog APIs to automatically log metrics, parameters, and fitted models to MLflow
  • (5 min) View results in the MLflow UI
  • (5 min) Use the MLflow Tracking APIs to track additional information (with example notebook)

Install MLflow

If you’re using Databricks Runtime for Machine Learning, MLflow is already installed. Otherwise, to install MLflow, add the following to the top of your notebook:

dbutils.library.installPyPI("mlflow")
dbutils.library.restartPython()
import mlflow

Automatically log training runs to MLflow

MLflow provides mlflow.<framework>.autolog() APIs to automatically log training code written in many ML frameworks. Call this API before running training code to log model-specific metrics, parameters, and model artifacts.

# Also autoinstruments tf.keras
import mlflow.tensorflow
mlflow.tensorflow.autolog()
# Use import mlflow.tensorflow and mlflow.tensorflow.autolog() if using tf.keras
import mlflow.keras
mlflow.keras.autolog()
import mlflow.xgboost
mlflow.xgboost.autolog()
import mlflow.lightgbm
mlflow.lightgbm.autolog()

If performing tuning with pyspark.ml, metrics and models are automatically logged to MLflow. See Apache Spark MLlib and automated MLflow tracking

View results

Once you’ve instrumented your machine learning code, you can view results in the Runs sidebar:

  1. Click the Runs icon in the notebook context bar. The Runs sidebar displays. In the sidebar, you can view the run’s parameters and metrics:

    Runs
  2. Click the External Link icon External Link in the Runs context bar to view the experiment:

    View experiment
  3. In the experiment, click a date:

    Select run

    The run details display:

    Run details
  4. In the experiment, click a source:

    Experiment source

    The notebook revision used in the run displays:

    Notebook revision

Track additional metrics, params, and models

You can log additional information by directly invoking the MLflow Tracking logging APIs.

  • Numerical metrics:

    import mlflow
    mlflow.log_metric("accuracy", 0.9)
    
  • Training parameters:

    import mlflow
    mlflow.log_param("learning_rate", 0.001)
    
  • Models:

    import mlflow.sklearn
    mlflow.sklearn.log_model(model, "myModel")
    
    import mlflow.spark
    mlflow.spark.log_model(model, "myModel")
    
    import mlflow.xgboost
    mlflow.xgboost.log_model(model, "myModel")
    
    import mlflow.keras
    mlflow.keras.log_model(model, "myModel")
    
    import mlflow.keras
    mlflow.keras.log_model(model, "myModel")
    
    import mlflow.pytorch
    mlflow.pytorch.log_model(model, "myModel")
    
    import mlflow.spacy
    mlflow.spacy.log_model(model, "myModel")
    
  • Other artifacts (files):

    import mlflow
    mlflow.log_artifact("/tmp/my-file", "myArtifactPath")
    

Example notebook

MLflow Quick Start Python notebook