Train And Deploy Keras Auto Logging

Copyright (c) Microsoft Corporation. All rights reserved.

Licensed under the MIT License.

Use MLflow with Azure Machine Learning to Train and Deploy Keras Image Classifier

This example shows you how to use MLflow together with Azure Machine Learning services for tracking the metrics and artifacts while training a Keras model to classify MNIST digit images and deploy the model as a web service. You'll learn how to:

1. Set up MLflow tracking URI so as to use Azure ML
2. Create experiment
3. Instrument your model with MLflow tracking
4. Train a Keras model locally with MLflow auto logging
5. Train a model on GPU compute on Azure with MLflow auto logging
6. View your experiment within your Azure ML Workspace in Azure Portal
7. Deploy the model as a web service on Azure Container Instance
8. Call the model to make predictions

Pre-requisites

If you are using a Notebook VM, you are all set. Otherwise, go through the Configuration notebook to set up your Azure Machine Learning workspace and ensure other common prerequisites are met.

Install TensorFlow and Keras, this notebook has been tested with TensorFlow version 2.1.0 and Keras version 2.3.1.

Also, install azureml-mlflow package using pip install azureml-mlflow. Note that azureml-mlflow installs mlflow package itself as a dependency if you haven't done so previously.

Set-up

Import packages and check versions of Azure ML SDK and MLflow installed on your computer. Then connect to your Workspace.

Set tracking URI

Set the MLflow tracking URI to point to your Azure ML Workspace. The subsequent logging calls from MLflow APIs will go to Azure ML services and will be tracked under your Workspace.

Create Experiment

In both MLflow and Azure ML, training runs are grouped into experiments. Let's create one for our experimentation.

Train model locally while logging metrics and artifacts

The scripts/train.py program contains the code to load the image dataset, train and test the model. Within this program, the train.driver function wraps the end-to-end workflow.

Within the driver, the mlflow.start_run starts MLflow tracking. Then, MLflow's automatic logging is used to log metrics, parameters and model for the Keras run.

Let's add the program to search path, import it as a module and invoke the driver function. Note that the training can take few minutes.

Train model on GPU compute on Azure

Next, let's run the same script on GPU-enabled compute for faster training. If you've completed the the Configuration notebook, you should have a GPU cluster named "gpu-cluster" available in your workspace. Otherwise, follow the instructions in the notebook to create one. For simplicity, this example uses single process on single VM to train the model.

Clone an environment object from the Tensorflow 2.1 Azure ML curated environment. Azure ML curated environments are pre-configured environments to simplify ML setup, reference this doc for more information. To enable MLflow tracking, add azureml-mlflow as pip package.

Create a ScriptRunConfig to specify the training configuration: script, compute as well as environment.

Get a reference to the experiment you created previously, but this time, as an Azure Machine Learning experiment object.

Then, use the Experiment.submit method to start the remote training run. Note that the first training run often takes longer as Azure Machine Learning service builds the Docker image for executing the script. Subsequent runs will be faster as the cached image is used.

You can monitor the run and its metrics on Azure Portal.

Also, you can wait for run to complete.

Deploy model as web service

The client.create_deployment function registers the logged Keras+Tensorflow model and deploys the model in a framework-aware manner. It automatically creates the Tensorflow-specific inferencing wrapper code and specifies package dependencies for you. See this doc for more information on deploying models on Azure ML using MLflow.

In this example, we deploy the Docker image to Azure Container Instance: a serverless compute capable of running a single container. You can tag and add descriptions to help keep track of your web service.

Other inferencing compute choices include Azure Kubernetes Service which provides scalable endpoint suitable for production use.

Note that the service deployment can take several minutes.

First define your deployment target and customize parameters in the deployment config. Refer to this documentation for more information.

Once the deployment has completed you can check the scoring URI of the web service in AzureML studio UI in the endpoints tab. Refer mlflow predict on how to test your deployment.

Clean up

You can delete the ACI deployment with a delete API call.