Automated MLflow tracking in MLlib

MLflow provides automated tracking for model tuning with MLlib. With automated MLflow tracking, when you run tuning code using CrossValidator or TrainValidationSplit, the specified hyperparameters and evaluation metrics are automatically logged, making it easy to identify the optimal model.

This notebook shows an example of automated MLflow tracking with MLlib.

This notebook uses the PySpark classes DecisionTreeClassifier and CrossValidator to train and tune a model. MLflow automatically tracks the learning process, saving the results of each run so you can examine the hyperparameters to understand the impact of each one on the model's performance and find the optimal settings.

This notebook uses the MNIST handwritten digit recognition dataset, which is included with Databricks.

Load the training and test datasets

The dataset is already divided into training and test sets. Each dataset has two columns: an image, represented as a vector of 784 pixels, and a "label", or the actual number shown in the image.

The datasets are stored in the LIBSVM dataset format. Load them using the MLlib LIBSVM dataset reader utility.

Display the data. Each image has the true label (the label column) and a vector of features that represent pixel intensities.

Define the ML pipeline

In this example, as with most ML applications, you must do some preprocessing of the data before you can use the data to train a model. MLlib provides pipelines that allow you to combine multiple steps into a single workflow. In this example, you build a two-step pipeline:

1. StringIndexer converts the labels from numeric values to categorical values.
2. DecisionTreeClassifier trains a decision tree that can predict the "label" column based on the data in the "features" column.

For more information:
Pipelines

Run the cross-validation

Now that you have defined the pipeline, you can run the cross validation to tune the model's hyperparameters. During this process, MLflow automatically tracks the models produced by CrossValidator, along with their evaluation metrics. This allows you to investigate how specific hyperparameters affect the model's performance.

In this example, you examine two hyperparameters in the cross-validation:

• maxDepth. This parameter determines how deep, and thus how large, the tree can grow.
• maxBins. For efficient distributed training of Decision Trees, MLlib discretizes (or "bins") continuous features into a finite number of values. The number of bins is controlled by maxBins. In this example, the number of bins corresponds to the number of grayscale levels; maxBins=2 turns the images into black and white images.

For more information:
maxBins
maxDepth

Run CrossValidator. If an MLflow tracking server is available, CrossValidator automatically logs each run to MLflow, along with the evaluation metric calculated on the held-out data, under the current active run. If no run is active, a new one is created.

Review the logged results

To view the MLflow experiment associated with the notebook, click the Experiment icon in the notebook context bar on the upper right. All notebook runs appear in the sidebar. To more easily compare their results, click the icon at the far right of Experiment Runs (it shows "View Experiment UI" when you hover over it). The Experiment page appears.

For example, to examine the effect of tuning maxDepth:

1. On the Experiment page, enter params.maxBins = "8" in the Search Runs box, and click Search.
2. Select the resulting runs and click Compare.
3. In the Scatter Plot, select X-axis maxDepth and Y-axis avg_weightedPrecision.

You can see that, when maxBins is held constant at 8, the average weighted precision increases with maxDepth.