Getting started with deep learning in Databricks: an end-to-end example using TensorFlow Keras, Hyperopt, and MLflow

This tutorial uses a small dataset to show how to use TensorFlow Keras, Hyperopt, and MLflow to develop a deep learning model in Databricks.

It includes the following steps:

• Load and preprocess data
• Part 1. Create a neural network model with TensorFlow Keras and view training with inline TensorBoard
• Part 2. Perform automated hyperparameter tuning with Hyperopt and MLflow and use autologging to save results
• Part 3. Use the best set of hyperparameters to build a final model
• Part 4. Register the model in MLflow and use the model to make predictions

Requirements

• Databricks Runtime for Machine Learning.

Load and preprocess data

This example uses the California Housing dataset from scikit-learn.

Scale features

Feature scaling is important when working with neural networks. This notebook uses the scikit-learn function StandardScaler.

Part 1. Create model and view TensorBoard in notebook

Create the neural network

Compile the model

Create callbacks

Start TensorBoard

TensorBoard continues to run until you detach the notebook from the cluster.
To clear the TensorBoard between runs, use this command: dbutils.fs.rm(experiment_log_dir.replace("/dbfs",""), recurse=True)

Evaluate model on test dataset

Part 2. Hyperparameter tuning with Hyperopt and MLflow

Hyperopt is a Python library for hyperparameter tuning. Databricks Runtime for Machine Learning includes an optimized and enhanced version of Hyperopt, including automated MLflow tracking. For more information about using Hyperopt, see the Hyperopt documentation.

Create neural network model using variables for number of nodes in hidden layers

Create Hyperopt objective function

Define Hyperopt search space

Create the SparkTrials object

The SparkTrials object tells fmin() to distribute the tuning job across a Spark cluster. When you create the SparkTrials object, you can use the parallelism argument to set the maximum number of trials to evaluate concurently. The default setting is the number of Spark executors available.

A higher number lets you scale-out testing of more hyperparameter settings. Because Hyperopt proposes new trials based on past results, there is a trade-off between parallelism and adaptivity. For a fixed max_evals, greater parallelism speeds up calculations, but lower parallelism may lead to better results since each iteration has access to more past results.

Perform hyperparameter tuning

Put the fmin() call inside an MLflow run to save results to MLflow. MLflow tracks the parameters and performance metrics of each run.

After running the following cell, you can view the results in MLflow. Click Experiment at the upper right to display the Experiment Runs sidebar. Click the icon at the far right next to Experiment Runs to display the MLflow Runs Table.

For more information about using MLflow to analyze runs, see (AWS|Azure|GCP).

Part 3. Use the best set of hyperparameters to build a final model

When autolog() is active, MLflow does not automatically end a run. You must end the run that was started in Cmd 30 before starting and autologging a new run.

Part 4. Register the model in MLflow and use the model to make predictions

To learn more about the Model Registry, see (AWS|Azure|GCP).

Load the model for inference and make predictions

Clean up

Detach this notebook from the cluster to stop TensorBoard.