Files
model_garden/notebooks/official/experiments/comparing_local_trained_models.ipynb
T
5b6c766629 Migrate gsutil usage to gcloud storage (#4298)
* Migrate gsutil usage to gcloud storage

* update

* linter changes for 4298

* Revert "linter changes for 4298"

This reverts commit c7a00a9710.

* Linter fixes

* update

* update

* Update distributed_hyperparameter_tuning.ipynb

* removed changes in model_garden folder

---------

Co-authored-by: gurusai-voleti <gvoleti@google.com>
2025-12-19 11:56:15 -05:00

26 KiB

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

Google Colaboratory logo
Open in Colab
Google Cloud Colab Enterprise logo
Open in Colab Enterprise
Vertex AI logo
Open in Workbench
GitHub logo
View on GitHub

Overview

As a Data Scientist, you may start running model experiments locally on your notebook. Depending on the framework you use, you need to track parameters, training time series and evaluation metrics. In this way, you are able to explain the modelling approach you have choosen.

Learn more about Vertex AI Experiments.

Objective

In this tutorial, you learn how to use Vertex AI Experiments to compare and evaluate model experiments.

This tutorial uses the following Google Cloud ML services and resources:

  • Vertex AI Workbench
  • Vertex AI Experiments

The steps performed include:

  • log the model parameters
  • log the loss and metrics on every epoch to TensorBoard
  • log the evaluation metrics

Dataset

In this notebook, you train a simple distributed neural network (DNN) model to predict an 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
  • Vertex AI TensorBoard
  • 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.

Get started

Install Vertex AI SDK for Python and other required packages

In [ ]:
! pip3 install --upgrade tensorflow==2.8 \
                        protobuf==3.20.3 \
                        google-cloud-aiplatform \
                        matplotlib \
                        pandas \
                        'numpy<2' -q --no-warn-conflicts

Restart runtime (Colab only)

To use the newly installed packages, you must restart the runtime on Google Colab.

In [ ]:
import sys

if "google.colab" in sys.modules:

    import IPython

    app = IPython.Application.instance()
    app.kernel.do_shutdown(True)
⚠️ The kernel is going to restart. Wait until it's finished before continuing to the next step. ⚠️

Authenticate your notebook environment (Colab only)

Authenticate your environment on Google Colab.

In [ ]:
import sys

if "google.colab" in sys.modules:

    from google.colab import auth

    auth.authenticate_user()

Set Google Cloud project information and initialize Vertex AI SDK for Python

To get started using Vertex AI, you must have an existing Google Cloud project and enable the Vertex AI API. Learn more about setting up a project and a development environment.

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

Create a Cloud Storage bucket

Create a storage bucket to store intermediate artifacts such as datasets.

In [ ]:
BUCKET_URI = f"gs://your-bucket-name-{PROJECT_ID}-unique"  # @param {type:"string"}

If your bucket doesn't already exist: Run the following cell to create your Cloud Storage bucket.

In [ ]:
! gcloud storage buckets create --location $LOCATION --project $PROJECT_ID $BUCKET_URI

Import libraries

In [ ]:
import uuid

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

Define constants

In [ ]:
EXPERIMENT_NAME = "[your-experiment-name]"  # @param {type:"string"}

If EXPERIMENT_NAME is not set, set a default one below:

In [ ]:
if EXPERIMENT_NAME == "[your-experiment-name]" or EXPERIMENT_NAME is None:
    EXPERIMENT_NAME = f"my-experiment-{uuid.uuid1()}"

Initialize Vertex AI SDK for Python

Initialize the Vertex AI SDK for Python for your project and corresponding bucket.

In [ ]:
vertex_ai.init(project=PROJECT_ID, location=LOCATION, staging_bucket=BUCKET_URI)

Create a Vertex AI Experiment with backing Vertex AI TensorBoard

You can create a Vertex AI Experiment by using the init() method. This automatically gets or creates the default Vertex AI TensorBoard and associates it with your Experiment.

Learn more about TensorBoard overview.

In [ ]:
vertex_ai.init(experiment=EXPERIMENT_NAME)

Training with Vertex AI Experiments

Vertex AI enables users to track the steps (for example, preprocessing, training) of an experiment run, and track inputs (for example, algorithm, parameters, datasets) and outputs (for example, models, checkpoints, metrics) of those steps.

To better understanding how parameters and metrics are stored and organized, the following concepts are explained:

  1. 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.

  2. 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.

You can use the Vertex AI SDK for Python to track metrics and parameters models trained locally for each experiment across several experiment runs.

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.

In [ ]:
# Helpers ----------------------------------------------------------------------


def read_data(uri):
    """
    Read data
    Args:
        uri: path to data
    Returns:
        pandas dataframe
    """
    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="", dtype=float)
    return dataset


def train_test_split(dataset, split_frac=0.8, random_state=0):
    """
    Split data into train and test
    Args:
        dataset: pandas dataframe
        split_frac: fraction of data to use for training
        random_state: random seed
    Returns:
        train and test dataframes
    """
    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


def normalize_dataset(train_dataset, test_dataset):
    """
    Normalize data
    Args:
        train_dataset: pandas dataframe
        test_dataset: pandas dataframe

    Returns:

    """
    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


def build_model(num_units, dropout_rate):
    """
    Build model
    Args:
        num_units: number of units in hidden layer
        dropout_rate: dropout rate
    Returns:
        compiled model
    """
    model = Sequential(
        [
            layers.Dense(
                num_units,
                activation="relu",
                input_shape=[9],
            ),
            layers.Dropout(rate=dropout_rate),
            layers.Dense(num_units, activation="relu"),
            layers.Dense(1),
        ]
    )

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


def train(
    model,
    train_data,
    train_labels,
    validation_split=0.2,
    epochs=10,
):
    """
    Train model
    Args:
        train_data: pandas dataframe
        train_labels: pandas dataframe
        model: compiled model
        validation_split: fraction of data to use for validation
        epochs: number of epochs to train for
    Returns:
        history
    """
    history = model.fit(
        train_data, train_labels, epochs=epochs, validation_split=validation_split
    )

    return history

Run experiment and evaluate experiment runs

You define several experiment configurations, run experiments and track them in Vertex AI Experiments

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

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

# Split data
train_dataset, test_dataset, train_labels, test_labels = train_test_split(dataset)

# Normalize data
normed_train_data, normed_test_data = normalize_dataset(train_dataset, test_dataset)

# Run experiments
for i, params in enumerate(parameters):

    # Initialize Vertex AI Experiment run
    vertex_ai.start_run(run=f"auto-mpg-local-run-{i}")

    # Log training parameters
    vertex_ai.log_params(params)

    # Build model
    model = build_model(
        num_units=params["num_units"], dropout_rate=params["dropout_rate"]
    )

    # Train model
    history = train(
        model,
        normed_train_data,
        train_labels,
        epochs=params["epochs"],
    )

    # Log additional parameters
    vertex_ai.log_params(history.params)

    # Log metrics per epochs
    for idx in range(0, history.params["epochs"]):
        vertex_ai.log_time_series_metrics(
            {
                "train_mae": history.history["mae"][idx],
                "train_mse": history.history["mse"][idx],
            }
        )

    # Log final metrics
    loss, mae, mse = model.evaluate(normed_test_data, test_labels, verbose=2)
    if np.isnan(loss):
        loss = 0
    if np.isnan(mae):
        mae = 0
    if np.isnan(mse):
        mse = 0
    vertex_ai.log_metrics({"eval_loss": loss, "eval_mae": mae, "eval_mse": mse})

    vertex_ai.end_run()

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 = vertex_ai.get_experiment_df()
experiment_df.T

Visualizing an experiment's parameters and metrics

You use parallel coordinates plotting to visualize 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.eval_loss",
        "metric.eval_mse",
        "metric.eval_mae",
    ],
    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.

In [ ]:
# Delete experiment
exp = vertex_ai.Experiment(EXPERIMENT_NAME)
backing_tensorboard = exp.get_backing_tensorboard_resource()
exp.delete(delete_backing_tensorboard_runs=True)

# Delete Tensorboard
delete_tensorboard = False  # Set True for deletion

if delete_tensorboard:
    backing_tensorboard.delete()

# Delete Cloud Storage objects that were created
delete_bucket = False  # Set True for deletion

if delete_bucket:
    ! gcloud storage rm --recursive --continue-on-error {BUCKET_URI}