model_garden/notebooks/official/ml_metadata/sdk-metric-parameter-tracking-for-locally-trained-models.ipynb

In [ ]:

# Copyright 2022 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

Vertex AI: Track parameters and metrics for locally trained models

Run in Colab

View on GitHub

Open in Vertex AI Workbench

Overview

This notebook demonstrates how to track metrics and parameters for ML training jobs and analyze this metadata using Vertex AI SDK for Python.

Learn more about Vertex ML Metadata

Objective

In this notebook, you learn how to use Vertex ML Metadata to track training parameters and evaluation metrics.

This tutorial uses the following Google Cloud ML services:

• Vertex ML Metadata
• Vertex AI Experiments

The steps performed include:

• Track parameters and metrics for a locally trained model.
• Extract and perform analysis for all parameters and metrics within an Experiment.

Dataset

In this notebook, we will train a simple distributed neural network (DNN) model to predict automobile's miles per gallon (MPG) based on automobile information in the auto-mpg dataset.

Costs

This tutorial uses billable components of Google Cloud:

• Vertex AI
• Cloud Storage

Learn about Vertex AI pricing and Cloud Storage pricing, and use the Pricing Calculator to generate a cost estimate based on your projected usage.

Installation

Install the packages required for executing this notebook.

In [ ]:

! pip install --upgrade --quiet google-cloud-aiplatform \
                                tensorflow==2.11

Colab only: Uncomment the following cell to restart the kernel

In [ ]:

# Automatically restart kernel after installs so that your environment can access the new packages
# import IPython

# app = IPython.Application.instance()
# app.kernel.do_shutdown(True)

Before you begin

Set your project ID

If you don't know your project ID, try the following:

• Run gcloud config list.
• Run gcloud project list.
• See the support page: Locate the project ID

In [ ]:

PROJECT_ID = "[your-project-id]"  # @param {type:"string"}

# Set the project id
! gcloud config set project $PROJECT_ID

Region

You can also change the REGION variable, used by Vertex AI. Learn more about Vertex AI regions

In [ ]:

REGION = "us-central1"  # @param {type: "string"}

Authenticate your Google Cloud account

Depending on your Jupyter environment, you may have to manually authenticate. Follow the relevant instructions below.

1. Vertex AI Workbench

• Do nothing as you are already authenticated.

2. Local JupyterLab instance, uncomment and run:

In [ ]:

# ! gcloud auth login

3. Colab, uncomment and run:

In [ ]:

# from google.colab import auth
# auth.authenticate_user()

4. Service Account or other

• See how to grant Cloud Storage permissions to your service account at https://cloud.google.com/storage/docs/gsutil/commands/iam#ch-examples

Import libraries and define constants

Import required libraries.

In [ ]:

import matplotlib.pyplot as plt
import pandas as pd
from google.cloud import aiplatform
from tensorflow.python.keras import Sequential, layers
from tensorflow.python.keras.utils import data_utils

Define some constants

In [ ]:

EXPERIMENT_NAME = "my-experiment-name-unique"  # @param {type:"string"}

Concepts

To better understanding how parameters and metrics are stored and organized, we'd like to introduce the following concepts:

Experiment

Experiments describe a context that groups your runs and the artifacts you create into a logical session. For example, in this notebook you create an Experiment and log data to that experiment.

Run

A run represents a single path/avenue that you executed while performing an experiment. A run includes artifacts that you used as inputs or outputs, and parameters that you used in this execution. An Experiment can contain multiple runs.

Getting started tracking parameters and metrics

You can use the Vertex SDK for Python to track metrics and parameters for models trained locally.

In the following example, you train a simple distributed neural network (DNN) model to predict automobile's miles per gallon (MPG) based on automobile information in the auto-mpg dataset.

Load and process the training dataset

Download and process the dataset.

In [ ]:

def read_data(uri):
    dataset_path = data_utils.get_file("auto-mpg.data", uri)
    column_names = [
        "MPG",
        "Cylinders",
        "Displacement",
        "Horsepower",
        "Weight",
        "Acceleration",
        "Model Year",
        "Origin",
    ]
    raw_dataset = pd.read_csv(
        dataset_path,
        names=column_names,
        na_values="?",
        comment="\t",
        sep=" ",
        skipinitialspace=True,
    )
    dataset = raw_dataset.dropna()
    dataset["Origin"] = dataset["Origin"].map(
        lambda x: {1: "USA", 2: "Europe", 3: "Japan"}.get(x)
    )
    dataset = pd.get_dummies(dataset, prefix="", prefix_sep="")
    return dataset


dataset = read_data(
    "http://archive.ics.uci.edu/ml/machine-learning-databases/auto-mpg/auto-mpg.data"
)

Split dataset for training and testing.

In [ ]:

def train_test_split(dataset, split_frac=0.8, random_state=0):
    train_dataset = dataset.sample(frac=split_frac, random_state=random_state)
    test_dataset = dataset.drop(train_dataset.index)
    train_labels = train_dataset.pop("MPG")
    test_labels = test_dataset.pop("MPG")

    return train_dataset, test_dataset, train_labels, test_labels


train_dataset, test_dataset, train_labels, test_labels = train_test_split(dataset)

Normalize the features in the dataset for better model performance.

In [ ]:

def normalize_dataset(train_dataset, test_dataset):
    train_stats = train_dataset.describe()
    train_stats = train_stats.transpose()

    def norm(x):
        return (x - train_stats["mean"]) / train_stats["std"]

    normed_train_data = norm(train_dataset)
    normed_test_data = norm(test_dataset)

    return normed_train_data, normed_test_data


normed_train_data, normed_test_data = normalize_dataset(train_dataset, test_dataset)

Define ML model and training function

In [ ]:

def train(
    train_data,
    train_labels,
    num_units=64,
    activation="relu",
    dropout_rate=0.0,
    validation_split=0.2,
    epochs=1000,
):

    model = Sequential(
        [
            layers.Dense(
                num_units,
                activation=activation,
                input_shape=[len(train_dataset.keys())],
            ),
            layers.Dropout(rate=dropout_rate),
            layers.Dense(num_units, activation=activation),
            layers.Dense(1),
        ]
    )

    model.compile(loss="mse", optimizer="adam", metrics=["mae", "mse"])
    print(model.summary())

    history = model.fit(
        train_data, train_labels, epochs=epochs, validation_split=validation_split
    )

    return model, history

Initialize the Vertex AI SDK for Python and create an Experiment

Initialize the client for Vertex AI and create an experiment.

In [ ]:

aiplatform.init(project=PROJECT_ID, location=REGION, experiment=EXPERIMENT_NAME)

Start several model training runs

Training parameters and metrics are logged for each run.

In [ ]:

parameters = [
    {"num_units": 16, "epochs": 3, "dropout_rate": 0.1},
    {"num_units": 16, "epochs": 10, "dropout_rate": 0.1},
    {"num_units": 16, "epochs": 10, "dropout_rate": 0.2},
    {"num_units": 32, "epochs": 10, "dropout_rate": 0.1},
    {"num_units": 32, "epochs": 10, "dropout_rate": 0.2},
]

for i, params in enumerate(parameters):
    aiplatform.start_run(run=f"auto-mpg-local-run-{i}")
    aiplatform.log_params(params)
    model, history = train(
        normed_train_data,
        train_labels,
        num_units=params["num_units"],
        activation="relu",
        epochs=params["epochs"],
        dropout_rate=params["dropout_rate"],
    )

    for metric, values in history.history.items():
        try:
            aiplatform.log_metrics({metric: values[-1]})
        except:
            aiplatform.log_metrics({metric: 0.0})

    loss, mae, mse = model.evaluate(normed_test_data, test_labels, verbose=2)
    try:
        aiplatform.log_metrics({"eval_loss": loss, "eval_mae": mae, "eval_mse": mse})
    except:
        aiplatform.log_metrics({"eval_loss": 0.0, "eval_mae": 0.0, "eval_mse": 0.0})

Extract parameters and metrics into a dataframe for analysis

We can also extract all parameters and metrics associated with any Experiment into a dataframe for further analysis.

In [ ]:

experiment_df = aiplatform.get_experiment_df()
experiment_df

Visualizing an experiment's parameters and metrics

In [ ]:

plt.rcParams["figure.figsize"] = [15, 5]

ax = pd.plotting.parallel_coordinates(
    experiment_df.reset_index(level=0),
    "run_name",
    cols=[
        "param.num_units",
        "param.dropout_rate",
        "param.epochs",
        "metric.loss",
        "metric.val_loss",
        "metric.eval_loss",
    ],
    color=["blue", "green", "pink", "red"],
)
ax.set_yscale("symlog")
ax.legend(bbox_to_anchor=(1.0, 0.5))

Visualizing experiments in Cloud Console

Run the following to get the URL of Vertex AI Experiments for your project.

In [ ]:

print("Vertex AI Experiments:")
print(
    f"https://console.cloud.google.com/ai/platform/experiments/experiments?folder=&organizationId=&project={PROJECT_ID}"
)

Cleaning up

To clean up all Google Cloud resources used in this project, you can delete the Google Cloud project you used for the tutorial.

Otherwise, you can delete the individual resources you created in this tutorial:

• Experiment (Can be deleted manually in the GCP Console UI)