ml-quickstart-training(Python)

Import Notebook

Databricks ML tutorial: Model Training

This notebook provides a quick overview of machine learning model training on Databricks. To train models, you can use libraries like scikit-learn that are preinstalled on the Databricks Runtime for Machine Learning. In addition, you can use MLflow to track the trained models, and Hyperopt with SparkTrials to scale hyperparameter tuning.

This tutorial covers:

  • Part 1: Training a simple classification model with MLflow tracking
  • Part 2: Hyperparameter tuning a better performing model with Hyperopt

For more details on productionizing machine learning on Databricks including model lifecycle management and model inference, see the ML End to End Example (AWS|Azure|GCP).

Requirements

  • Cluster running Databricks Runtime 15.4 LTS ML or above

Libraries

Import the necessary libraries. These libraries are preinstalled on Databricks Runtime for Machine Learning (AWS|Azure|GCP) clusters and are tuned for compatibility and performance.

3

Load data

The tutorial uses a dataset describing different wine samples. The dataset is from the UCI Machine Learning Repository. The goal is to classify red and white wines by their quality.

For more details on uploading and loading from other data sources, see the documentation on working with data (AWS|Azure|GCP).

5

Part 1. Train a classification model

MLflow Tracking

MLflow tracking allows you to organize your machine learning training code, parameters, and models.

You can enable automatic MLflow tracking by using autologging.

8

2024/10/22 21:21:58 INFO mlflow.tracking.fluent: Autologging successfully enabled for sklearn. 2024/10/22 21:21:58 WARNING mlflow.spark: With Pyspark >= 3.2, PYSPARK_PIN_THREAD environment variable must be set to false for Spark datasource autologging to work. 2024/10/22 21:21:58 INFO mlflow.tracking.fluent: Autologging successfully enabled for pyspark. 2024/10/22 21:21:58 INFO mlflow.tracking.fluent: Autologging successfully enabled for pyspark.ml.

Next, train a classifier within the context of an MLflow run, which automatically logs the trained model and many associated metrics and parameters.

You can supplement the logging with additional metrics such as the model's AUC score on the test dataset.

10

2024/10/22 21:22:05 WARNING mlflow.utils.autologging_utils: MLflow autologging encountered a warning: "/databricks/python/lib/python3.11/site-packages/_distutils_hack/__init__.py:33: UserWarning: Setuptools is replacing distutils."
Test AUC of: 0.8834365701533531 2024/10/22 21:22:08 INFO mlflow.tracking._tracking_service.client: 🏃 View run gradient_boost at: db-sme-demo-docs.cloud.databricks.com/ml/experiments/3770961151225428/runs/be6acf40f819432788b4c9746d04524e. 2024/10/22 21:22:08 INFO mlflow.tracking._tracking_service.client: 🧪 View experiment at: db-sme-demo-docs.cloud.databricks.com/ml/experiments/3770961151225428.

If you aren't happy with the performance of this model, train another model with different hyperparameters.

12

View MLflow runs

To view the logged training runs, click the Experiment icon at the upper right of the notebook to display the experiment sidebar. If necessary, click the refresh icon to fetch and monitor the latest runs.

You can then click the experiment page icon to display the more detailed MLflow experiment page (AWS|Azure|GCP). This page allows you to compare runs and view details for specific runs.

Load models

You can also access the results for a specific run using the MLflow API. The code in the following cell illustrates how to load the model trained in a given MLflow run and use it to make predictions. You can also find code snippets for loading specific models on the MLflow run page (AWS|Azure|GCP).

15

Part 2. Hyperparameter Tuning

At this point, you have trained a simple model and used the MLflow tracking service to organize your work. This section covers how to perform more sophisticated tuning using Hyperopt.

Parallel training with Hyperopt and SparkTrials

Hyperopt is a Python library for hyperparameter tuning. For more information about using Hyperopt in Databricks, see the documentation (AWS|Azure|GCP).

You can use Hyperopt with SparkTrials to run hyperparameter sweeps and train multiple models in parallel. This reduces the time required to optimize model performance. MLflow tracking is integrated with Hyperopt to automatically log models and parameters.

18

Hyperopt with SparkTrials will automatically track trials in MLflow. To view the MLflow experiment associated with the notebook, click the 'Runs' icon in the notebook context bar on the upper right. There, you can view all runs. To view logs from trials, please check the Spark executor logs. To view executor logs, expand 'Spark Jobs' above until you see the (i) icon next to the stage from the trial job. Click it and find the list of tasks. Click the 'stderr' link for a task to view trial logs. 0%| | 0/32 [00:00<?, ?trial/s, best loss=?]2024/10/22 21:22:43 INFO mlflow.tracking._tracking_service.client: 🏃 View run big-pig-212 at: db-sme-demo-docs.cloud.databricks.com/ml/experiments/3770961151225428/runs/7d686e19f1d04a86acd7a5033fca9af8. 2024/10/22 21:22:43 INFO mlflow.tracking._tracking_service.client: 🧪 View experiment at: db-sme-demo-docs.cloud.databricks.com/ml/experiments/3770961151225428. 3%|▎ | 1/32 [00:21<11:00, 21.30s/trial, best loss: -0.8991760128175784]2024/10/22 21:22:48 INFO mlflow.tracking._tracking_service.client: 🏃 View run grandiose-quail-616 at: db-sme-demo-docs.cloud.databricks.com/ml/experiments/3770961151225428/runs/6cc4339652e64aa8adadee6f2d59d046. 2024/10/22 21:22:48 INFO mlflow.tracking._tracking_service.client: 🧪 View experiment at: db-sme-demo-docs.cloud.databricks.com/ml/experiments/3770961151225428. 6%|▋ | 2/32 [00:25<05:33, 11.13s/trial, best loss: -0.8991760128175784]2024/10/22 21:23:05 INFO mlflow.tracking._tracking_service.client: 🏃 View run carefree-shoat-681 at: db-sme-demo-docs.cloud.databricks.com/ml/experiments/3770961151225428/runs/2cf5fbbcdaac41b59eb06ff0548e82e9. 2024/10/22 21:23:05 INFO mlflow.tracking._tracking_service.client: 🧪 View experiment at: db-sme-demo-docs.cloud.databricks.com/ml/experiments/3770961151225428. 12%|█▎ | 4/32 [00:56<06:25, 13.77s/trial, best loss: -0.8991760128175784]2024/10/22 21:23:19 INFO mlflow.tracking._tracking_service.client: 🏃 View run bright-dog-904 at: db-sme-demo-docs.cloud.databricks.com/ml/experiments/3770961151225428/runs/af76b3c81da84dc0b8fe684478c14550. 2024/10/22 21:23:19 INFO mlflow.tracking._tracking_service.client: 🧪 View experiment at: db-sme-demo-docs.cloud.databricks.com/ml/experiments/3770961151225428. 16%|█▌ | 5/32 [01:17<07:24, 16.45s/trial, best loss: -0.9096260585946442]2024/10/22 21:23:40 INFO mlflow.tracking._tracking_service.client: 🏃 View run secretive-turtle-294 at: db-sme-demo-docs.cloud.databricks.com/ml/experiments/3770961151225428/runs/c3c128755ff5455caeb55fa52310ada3. 2024/10/22 21:23:40 INFO mlflow.tracking._tracking_service.client: 🧪 View experiment at: db-sme-demo-docs.cloud.databricks.com/ml/experiments/3770961151225428. 19%|█▉ | 6/32 [01:31<06:46, 15.62s/trial, best loss: -0.9096260585946442]2024/10/22 21:23:54 INFO mlflow.tracking._tracking_service.client: 🏃 View run painted-mare-353 at: db-sme-demo-docs.cloud.databricks.com/ml/experiments/3770961151225428/runs/905214357a7941f791a89ab16a8ed09f. 2024/10/22 21:23:54 INFO mlflow.tracking._tracking_service.client: 🧪 View experiment at: db-sme-demo-docs.cloud.databricks.com/ml/experiments/3770961151225428. 22%|██▏ | 7/32 [01:45<06:18, 15.16s/trial, best loss: -0.9148260471503777]2024/10/22 21:24:08 INFO mlflow.tracking._tracking_service.client: 🏃 View run efficient-croc-93 at: db-sme-demo-docs.cloud.databricks.com/ml/experiments/3770961151225428/runs/93613c8cd8ca46498118ac3dc6638f20. 2024/10/22 21:24:08 INFO mlflow.tracking._tracking_service.client: 🧪 View experiment at: db-sme-demo-docs.cloud.databricks.com/ml/experiments/3770961151225428. 25%|██▌ | 8/32 [02:04<06:33, 16.38s/trial, best loss: -0.9148260471503777]2024/10/22 21:24:28 INFO mlflow.tracking._tracking_service.client: 🏃 View run nimble-seal-488 at: db-sme-demo-docs.cloud.databricks.com/ml/experiments/3770961151225428/runs/abe20f8466634911b892ba1c75a81a72. 2024/10/22 21:24:28 INFO mlflow.tracking._tracking_service.client: 🧪 View experiment at: db-sme-demo-docs.cloud.databricks.com/ml/experiments/3770961151225428. 2024/10/22 21:24:44 INFO mlflow.tracking._tracking_service.client: 🏃 View run flawless-smelt-963 at: db-sme-demo-docs.cloud.databricks.com/ml/experiments/3770961151225428/runs/0df6a4ae56034fe99ed9ee478c85036a. 2024/10/22 21:24:44 INFO mlflow.tracking._tracking_service.client: 🧪 View experiment at: db-sme-demo-docs.cloud.databricks.com/ml/experiments/3770961151225428. 28%|██▊ | 9/32 [02:21<06:21, 16.58s/trial, best loss: -0.9148260471503777]2024/10/22 21:24:54 INFO mlflow.tracking._tracking_service.client: 🏃 View run indecisive-hound-317 at: db-sme-demo-docs.cloud.databricks.com/ml/experiments/3770961151225428/runs/8d1cfd68266e4ed1b0ff0402e93d7f93. 2024/10/22 21:24:54 INFO mlflow.tracking._tracking_service.client: 🧪 View experiment at: db-sme-demo-docs.cloud.databricks.com/ml/experiments/3770961151225428. 31%|███▏ | 10/32 [02:31<05:20, 14.55s/trial, best loss: -0.9148260471503777]2024/10/22 21:25:10 INFO mlflow.tracking._tracking_service.client: 🏃 View run languid-wren-177 at: db-sme-demo-docs.cloud.databricks.com/ml/experiments/3770961151225428/runs/55e59bb946634feb971dc6c56825653a. 2024/10/22 21:25:10 INFO mlflow.tracking._tracking_service.client: 🧪 View experiment at: db-sme-demo-docs.cloud.databricks.com/ml/experiments/3770961151225428. 38%|███▊ | 12/32 [02:57<04:29, 13.48s/trial, best loss: -0.9148260471503777]2024/10/22 21:25:21 INFO mlflow.tracking._tracking_service.client: 🏃 View run zealous-fox-637 at: db-sme-demo-docs.cloud.databricks.com/ml/experiments/3770961151225428/runs/2179eb34b02f4e5cb999c2d21efa7360. 2024/10/22 21:25:21 INFO mlflow.tracking._tracking_service.client: 🧪 View experiment at: db-sme-demo-docs.cloud.databricks.com/ml/experiments/3770961151225428. 41%|████ | 13/32 [03:08<04:01, 12.73s/trial, best loss: -0.9148260471503777]2024/10/22 21:25:32 INFO mlflow.tracking._tracking_service.client: 🏃 View run nosy-wasp-493 at: db-sme-demo-docs.cloud.databricks.com/ml/experiments/3770961151225428/runs/7fb6e830cfa0460e88aed6339fc3b1f0. 2024/10/22 21:25:32 INFO mlflow.tracking._tracking_service.client: 🧪 View experiment at: db-sme-demo-docs.cloud.databricks.com/ml/experiments/3770961151225428. 44%|████▍ | 14/32 [03:26<04:17, 14.33s/trial, best loss: -0.9148260471503777]2024/10/22 21:25:49 INFO mlflow.tracking._tracking_service.client: 🏃 View run invincible-skunk-277 at: db-sme-demo-docs.cloud.databricks.com/ml/experiments/3770961151225428/runs/d3e96d2671734b48b1e35c19bd92f47b. 2024/10/22 21:25:49 INFO mlflow.tracking._tracking_service.client: 🧪 View experiment at: db-sme-demo-docs.cloud.databricks.com/ml/experiments/3770961151225428. 2024/10/22 21:26:06 INFO mlflow.tracking._tracking_service.client: 🏃 View run gaudy-cod-744 at: db-sme-demo-docs.cloud.databricks.com/ml/experiments/3770961151225428/runs/a6b8ad2a3ba94175b2f3df9b9b5c5e05. 2024/10/22 21:26:06 INFO mlflow.tracking._tracking_service.client: 🧪 View experiment at: db-sme-demo-docs.cloud.databricks.com/ml/experiments/3770961151225428. 50%|█████ | 16/32 [03:53<03:37, 13.59s/trial, best loss: -0.9148260471503777]2024/10/22 21:26:16 INFO mlflow.tracking._tracking_service.client: 🏃 View run omniscient-flea-793 at: db-sme-demo-docs.cloud.databricks.com/ml/experiments/3770961151225428/runs/5e02e335fc9548feb3e9f6f4d188571c. 2024/10/22 21:26:16 INFO mlflow.tracking._tracking_service.client: 🧪 View experiment at: db-sme-demo-docs.cloud.databricks.com/ml/experiments/3770961151225428. 53%|█████▎ | 17/32 [04:12<03:48, 15.22s/trial, best loss: -0.9148260471503777]2024/10/22 21:26:35 INFO mlflow.tracking._tracking_service.client: 🏃 View run enchanting-quail-502 at: db-sme-demo-docs.cloud.databricks.com/ml/experiments/3770961151225428/runs/b4a7643dd7f341248b56acbbb67c0a09. 2024/10/22 21:26:35 INFO mlflow.tracking._tracking_service.client: 🧪 View experiment at: db-sme-demo-docs.cloud.databricks.com/ml/experiments/3770961151225428. 56%|█████▋ | 18/32 [04:28<03:36, 15.46s/trial, best loss: -0.9148260471503777]2024/10/22 21:26:52 INFO mlflow.tracking._tracking_service.client: 🏃 View run clumsy-hawk-198 at: db-sme-demo-docs.cloud.databricks.com/ml/experiments/3770961151225428/runs/3df8b228a9c847b2a8fdaeeb6751b71f. 2024/10/22 21:26:52 INFO mlflow.tracking._tracking_service.client: 🧪 View experiment at: db-sme-demo-docs.cloud.databricks.com/ml/experiments/3770961151225428. 59%|█████▉ | 19/32 [04:51<03:47, 17.50s/trial, best loss: -0.9164210917830167]2024/10/22 21:27:14 INFO mlflow.tracking._tracking_service.client: 🏃 View run polite-cub-372 at: db-sme-demo-docs.cloud.databricks.com/ml/experiments/3770961151225428/runs/e208b4e5e87049e7bad4f38b8eaf55a6. 2024/10/22 21:27:14 INFO mlflow.tracking._tracking_service.client: 🧪 View experiment at: db-sme-demo-docs.cloud.databricks.com/ml/experiments/3770961151225428. 2024/10/22 21:27:24 INFO mlflow.tracking._tracking_service.client: 🏃 View run shivering-gull-679 at: db-sme-demo-docs.cloud.databricks.com/ml/experiments/3770961151225428/runs/f22860fbd51c4f49ae06e27d49566637. 2024/10/22 21:27:24 INFO mlflow.tracking._tracking_service.client: 🧪 View experiment at: db-sme-demo-docs.cloud.databricks.com/ml/experiments/3770961151225428. 66%|██████▌ | 21/32 [05:16<02:46, 15.09s/trial, best loss: -0.9164210917830167]2024/10/22 21:27:39 INFO mlflow.tracking._tracking_service.client: 🏃 View run treasured-shrew-397 at: db-sme-demo-docs.cloud.databricks.com/ml/experiments/3770961151225428/runs/658e4387bd2745fab57ec8d59346dcb3. 2024/10/22 21:27:39 INFO mlflow.tracking._tracking_service.client: 🧪 View experiment at: db-sme-demo-docs.cloud.databricks.com/ml/experiments/3770961151225428. 2024/10/22 21:27:58 INFO mlflow.tracking._tracking_service.client: 🏃 View run luxuriant-carp-75 at: db-sme-demo-docs.cloud.databricks.com/ml/experiments/3770961151225428/runs/f42b473265ec47eab4e6e449a3d9bd1d. 2024/10/22 21:27:58 INFO mlflow.tracking._tracking_service.client: 🧪 View experiment at: db-sme-demo-docs.cloud.databricks.com/ml/experiments/3770961151225428. 69%|██████▉ | 22/32 [05:36<02:45, 16.57s/trial, best loss: -0.9164210917830167]2024/10/22 21:28:08 INFO mlflow.tracking._tracking_service.client: 🏃 View run exultant-toad-793 at: db-sme-demo-docs.cloud.databricks.com/ml/experiments/3770961151225428/runs/1b885657f7c444539801ac7deaba643a. 2024/10/22 21:28:08 INFO mlflow.tracking._tracking_service.client: 🧪 View experiment at: db-sme-demo-docs.cloud.databricks.com/ml/experiments/3770961151225428. 75%|███████▌ | 24/32 [06:00<01:55, 14.42s/trial, best loss: -0.9164210917830167]2024/10/22 21:28:23 INFO mlflow.tracking._tracking_service.client: 🏃 View run big-foal-62 at: db-sme-demo-docs.cloud.databricks.com/ml/experiments/3770961151225428/runs/bbb12b3fe63847c5948c456ba0886976. 2024/10/22 21:28:23 INFO mlflow.tracking._tracking_service.client: 🧪 View experiment at: db-sme-demo-docs.cloud.databricks.com/ml/experiments/3770961151225428. 78%|███████▊ | 25/32 [06:18<01:48, 15.50s/trial, best loss: -0.9164210917830167]2024/10/22 21:28:41 INFO mlflow.tracking._tracking_service.client: 🏃 View run dazzling-mule-422 at: db-sme-demo-docs.cloud.databricks.com/ml/experiments/3770961151225428/runs/d7a6ef4fcc674d6296df7cbad2b9f5d6. 2024/10/22 21:28:41 INFO mlflow.tracking._tracking_service.client: 🧪 View experiment at: db-sme-demo-docs.cloud.databricks.com/ml/experiments/3770961151225428. 2024/10/22 21:28:54 INFO mlflow.tracking._tracking_service.client: 🏃 View run gaudy-cat-796 at: db-sme-demo-docs.cloud.databricks.com/ml/experiments/3770961151225428/runs/c7bc172cbc554e1699d435c72ef63607. 2024/10/22 21:28:54 INFO mlflow.tracking._tracking_service.client: 🧪 View experiment at: db-sme-demo-docs.cloud.databricks.com/ml/experiments/3770961151225428. 84%|████████▍ | 27/32 [06:46<01:13, 14.74s/trial, best loss: -0.9164210917830167]2024/10/22 21:29:09 INFO mlflow.tracking._tracking_service.client: 🏃 View run loud-cow-428 at: db-sme-demo-docs.cloud.databricks.com/ml/experiments/3770961151225428/runs/148c736de0264179ab4b471cbf482e01. 2024/10/22 21:29:09 INFO mlflow.tracking._tracking_service.client: 🧪 View experiment at: db-sme-demo-docs.cloud.databricks.com/ml/experiments/3770961151225428. 88%|████████▊ | 28/32 [07:07<01:06, 16.62s/trial, best loss: -0.9164210917830167]2024/10/22 21:29:30 INFO mlflow.tracking._tracking_service.client: 🏃 View run abundant-boar-672 at: db-sme-demo-docs.cloud.databricks.com/ml/experiments/3770961151225428/runs/28fcaceb2f594d3f9b0d26d532878eb0. 2024/10/22 21:29:30 INFO mlflow.tracking._tracking_service.client: 🧪 View experiment at: db-sme-demo-docs.cloud.databricks.com/ml/experiments/3770961151225428. 91%|█████████ | 29/32 [07:18<00:44, 14.93s/trial, best loss: -0.9164210917830167]2024/10/22 21:29:41 INFO mlflow.tracking._tracking_service.client: 🏃 View run caring-snake-246 at: db-sme-demo-docs.cloud.databricks.com/ml/experiments/3770961151225428/runs/920791729296491b929797f1993591a3. 2024/10/22 21:29:41 INFO mlflow.tracking._tracking_service.client: 🧪 View experiment at: db-sme-demo-docs.cloud.databricks.com/ml/experiments/3770961151225428. 2024/10/22 21:29:56 INFO mlflow.tracking._tracking_service.client: 🏃 View run enchanting-eel-15 at: db-sme-demo-docs.cloud.databricks.com/ml/experiments/3770961151225428/runs/489cc862c3834eb29caafd519fdb0019. 2024/10/22 21:29:56 INFO mlflow.tracking._tracking_service.client: 🧪 View experiment at: db-sme-demo-docs.cloud.databricks.com/ml/experiments/3770961151225428. 97%|█████████▋| 31/32 [07:52<00:16, 16.08s/trial, best loss: -0.9164210917830167]2024/10/22 21:30:15 INFO mlflow.tracking._tracking_service.client: 🏃 View run intrigued-bat-650 at: db-sme-demo-docs.cloud.databricks.com/ml/experiments/3770961151225428/runs/7755d1b7cd2244e281982f05256dd9dc. 2024/10/22 21:30:15 INFO mlflow.tracking._tracking_service.client: 🧪 View experiment at: db-sme-demo-docs.cloud.databricks.com/ml/experiments/3770961151225428. 100%|██████████| 32/32 [08:03<00:00, 15.10s/trial, best loss: -0.9164210917830167] 2024/10/22 21:30:26 INFO mlflow.tracking._tracking_service.client: 🏃 View run skillful-gnat-457 at: db-sme-demo-docs.cloud.databricks.com/ml/experiments/3770961151225428/runs/c7da330bf0f54fab9e355cbe244e9911. 2024/10/22 21:30:26 INFO mlflow.tracking._tracking_service.client: 🧪 View experiment at: db-sme-demo-docs.cloud.databricks.com/ml/experiments/3770961151225428. Total Trials: 32: 32 succeeded, 0 failed, 0 cancelled. 2024/10/22 21:30:27 INFO mlflow.tracking._tracking_service.client: 🏃 View run gb_hyperopt at: db-sme-demo-docs.cloud.databricks.com/ml/experiments/3770961151225428/runs/3f1dcc8cfc03470f97bfc1e13192fa08. 2024/10/22 21:30:27 INFO mlflow.tracking._tracking_service.client: 🧪 View experiment at: db-sme-demo-docs.cloud.databricks.com/ml/experiments/3770961151225428.

Search runs to retrieve the best model

Because all of the runs are tracked by MLflow, you can retrieve the metrics and parameters for the best run using the MLflow search runs API to find the tuning run with the highest test auc.

This tuned model should perform better than the simpler models trained in Part 1.

20

Best Run AUC: 0.9164210917830167 Num Estimators: 868.0 Max Depth: 5.0 Learning Rate: 0.10093731909149309
Test Predictions: [False False False True False]