fix: hanging on testing

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

@@ -0,0 +1,890 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "id": "ur8xi4C7S06n"
+   },
+   "outputs": [],
+   "source": [
+    "# Copyright 2021 Google LLC\n",
+    "#\n",
+    "# Licensed under the Apache License, Version 2.0 (the \"License\");\n",
+    "# you may not use this file except in compliance with the License.\n",
+    "# You may obtain a copy of the License at\n",
+    "#\n",
+    "#     https://www.apache.org/licenses/LICENSE-2.0\n",
+    "#\n",
+    "# Unless required by applicable law or agreed to in writing, software\n",
+    "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
+    "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
+    "# See the License for the specific language governing permissions and\n",
+    "# limitations under the License."
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "id": "JAPoU8Sm5E6e"
+   },
+   "source": [
+    "<table align=\"left\">\n",
+    "\n",
+    "  <td>\n",
+    "    <a href=\"https://colab.research.google.com/github/GoogleCloudPlatform/vertex-ai-samples/blob/master/notebooks/official/ml_metadata/sdk-metric-parameter-tracking-for-locally-trained-models.ipynb\">\n",
+    "      <img src=\"https://cloud.google.com/ml-engine/images/colab-logo-32px.png\" alt=\"Colab logo\"> Run in Colab\n",
+    "    </a>\n",
+    "  </td>\n",
+    "  <td>\n",
+    "    <a href=\"https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/master/notebooks/official/ml_metadata/sdk-metric-parameter-tracking-for-locally-trained-models.ipynb\">\n",
+    "      <img src=\"https://cloud.google.com/ml-engine/images/github-logo-32px.png\" alt=\"GitHub logo\">\n",
+    "      View on GitHub\n",
+    "    </a>\n",
+    "  </td>\n",
+    "</table>"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "id": "WBFL9LagqmwT"
+   },
+   "source": [
+    "#Vertex AI: Track parameters and metrics for locally trained models"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "id": "tvgnzT1CKxrO"
+   },
+   "source": [
+    "## Overview\n",
+    "\n",
+    "This notebook demonstrates how to track metrics and parameters for ML training jobs and analyze this metadata using Vertex AI SDK.\n",
+    "\n",
+    "### Dataset\n",
+    "\n",
+    "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](https://www.kaggle.com/devanshbesain/exploration-and-analysis-auto-mpg).\n",
+    "\n",
+    "### Objective\n",
+    "\n",
+    "In this notebook, you will learn how to use Vertex AI SDK to:\n",
+    "\n",
+    "    * Track parameters and metrics for a locally trainined model.\n",
+    "    * Extract and perform analysis for all parameters and metrics within an Experiment.\n",
+    "\n",
+    "### Costs \n",
+    "\n",
+    "\n",
+    "This tutorial uses billable components of Google Cloud:\n",
+    "\n",
+    "* Vertex AI\n",
+    "* Cloud Storage\n",
+    "\n",
+    "\n",
+    "Learn about [Vertex AI\n",
+    "pricing](https://cloud.google.com/vertex-ai/pricing) and [Cloud Storage\n",
+    "pricing](https://cloud.google.com/storage/pricing), and use the [Pricing\n",
+    "Calculator](https://cloud.google.com/products/calculator/)\n",
+    "to generate a cost estimate based on your projected usage."
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "id": "ze4-nDLfK4pw"
+   },
+   "source": [
+    "### Set up your local development environment\n",
+    "\n",
+    "**If you are using Colab or AI Platform Notebooks**, your environment already meets\n",
+    "all the requirements to run this notebook. You can skip this step."
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "id": "gCuSR8GkAgzl"
+   },
+   "source": [
+    "**Otherwise**, make sure your environment meets this notebook's requirements.\n",
+    "You need the following:\n",
+    "\n",
+    "* The Google Cloud SDK\n",
+    "* Git\n",
+    "* Python 3\n",
+    "* virtualenv\n",
+    "* Jupyter notebook running in a virtual environment with Python 3\n",
+    "\n",
+    "The Google Cloud guide to [Setting up a Python development\n",
+    "environment](https://cloud.google.com/python/setup) and the [Jupyter\n",
+    "installation guide](https://jupyter.org/install) provide detailed instructions\n",
+    "for meeting these requirements. The following steps provide a condensed set of\n",
+    "instructions:\n",
+    "\n",
+    "1. [Install and initialize the Cloud SDK.](https://cloud.google.com/sdk/docs/)\n",
+    "\n",
+    "1. [Install Python 3.](https://cloud.google.com/python/setup#installing_python)\n",
+    "\n",
+    "1. [Install\n",
+    "   virtualenv](https://cloud.google.com/python/setup#installing_and_using_virtualenv)\n",
+    "   and create a virtual environment that uses Python 3. Activate the virtual environment.\n",
+    "\n",
+    "1. To install Jupyter, run `pip3 install jupyter` on the\n",
+    "command-line in a terminal shell.\n",
+    "\n",
+    "1. To launch Jupyter, run `jupyter notebook` on the command-line in a terminal shell.\n",
+    "\n",
+    "1. Open this notebook in the Jupyter Notebook Dashboard."
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "id": "i7EUnXsZhAGF"
+   },
+   "source": [
+    "### Install additional packages\n",
+    "\n",
+    "Run the following commands to install the Vertex AI SDK and packages used in this notebook."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "id": "IaYsrh0Tc17L"
+   },
+   "outputs": [],
+   "source": [
+    "import os\n",
+    "\n",
+    "# The Google Cloud Notebook product has specific requirements\n",
+    "IS_GOOGLE_CLOUD_NOTEBOOK = os.path.exists(\"/opt/deeplearning/metadata/env_version\")\n",
+    "\n",
+    "# Google Cloud Notebook requires dependencies to be installed with '--user'\n",
+    "USER_FLAG = \"\"\n",
+    "if IS_GOOGLE_CLOUD_NOTEBOOK:\n",
+    "    USER_FLAG = \"--user\""
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "id": "MCQDRsnE3uzz"
+   },
+   "source": [
+    "Install Vertex AI SDK, matplotlib and tensorflow for training and evaluation model."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "id": "wyy5Lbnzg5fi"
+   },
+   "outputs": [],
+   "source": [
+    "! pip3 install {USER_FLAG} --upgrade tensorflow\n",
+    "! pip3 install {USER_FLAG} --upgrade google-cloud-aiplatform"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "id": "hhq5zEbGg0XX"
+   },
+   "source": [
+    "### Restart the kernel\n",
+    "\n",
+    "After you install the additional packages, you need to restart the notebook kernel so it can find the packages."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "id": "EzrelQZ22IZj"
+   },
+   "outputs": [],
+   "source": [
+    "# Automatically restart kernel after installs\n",
+    "import os\n",
+    "\n",
+    "if not os.getenv(\"IS_TESTING\"):\n",
+    "    # Automatically restart kernel after installs\n",
+    "    import IPython\n",
+    "\n",
+    "    app = IPython.Application.instance()\n",
+    "    app.kernel.do_shutdown(True)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "id": "lWEdiXsJg0XY"
+   },
+   "source": [
+    "## Before you begin\n",
+    "\n",
+    "### Select a GPU runtime\n",
+    "\n",
+    "**Make sure you're running this notebook in a GPU runtime if you have that option. In Colab, select \"Runtime --> Change runtime type > GPU\"**"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "id": "BF1j6f9HApxa"
+   },
+   "source": [
+    "### Set up your Google Cloud project\n",
+    "\n",
+    "**The following steps are required, regardless of your notebook environment.**\n",
+    "\n",
+    "1. [Select or create a Google Cloud project](https://console.cloud.google.com/cloud-resource-manager). When you first create an account, you get a $300 free credit towards your compute/storage costs.\n",
+    "\n",
+    "1. [Make sure that billing is enabled for your project](https://cloud.google.com/billing/docs/how-to/modify-project).\n",
+    "\n",
+    "1. [Enable the Vertex AI API](https://console.cloud.google.com/flows/enableapi?apiid=aiplatform.googleapis.com).\n",
+    "\n",
+    "1. If you are running this notebook locally, you will need to install the [Cloud SDK](https://cloud.google.com/sdk).\n",
+    "\n",
+    "1. Enter your project ID in the cell below. Then run the cell to make sure the\n",
+    "Cloud SDK uses the right project for all the commands in this notebook.\n",
+    "\n",
+    "**Note**: Jupyter runs lines prefixed with `!` as shell commands, and it interpolates Python variables prefixed with `$` into these commands."
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "id": "WReHDGG5g0XY"
+   },
+   "source": [
+    "#### Set your project ID\n",
+    "\n",
+    "**If you don't know your project ID**, you may be able to get your project ID using `gcloud`."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "id": "oM1iC_MfAts1"
+   },
+   "outputs": [],
+   "source": [
+    "import os\n",
+    "\n",
+    "PROJECT_ID = \"\"\n",
+    "\n",
+    "# Get your Google Cloud project ID from gcloud\n",
+    "if not os.getenv(\"IS_TESTING\"):\n",
+    "    shell_output=!gcloud config list --format 'value(core.project)' 2>/dev/null\n",
+    "    PROJECT_ID = shell_output[0]\n",
+    "    print(\"Project ID: \", PROJECT_ID)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "id": "qJYoRfYng0XZ"
+   },
+   "source": [
+    "Otherwise, set your project ID here."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "id": "riG_qUokg0XZ"
+   },
+   "outputs": [],
+   "source": [
+    "if PROJECT_ID == \"\" or PROJECT_ID is None:\n",
+    "    PROJECT_ID = \"[your-project-id]\"  # @param {type:\"string\"}"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "id": "06571eb4063b"
+   },
+   "source": [
+    "#### Timestamp\n",
+    "\n",
+    "If you are in a live tutorial session, you might be using a shared test account or project. To avoid name collisions between users on resources created, you create a timestamp for each instance session, and append it onto the name of resources you create in this tutorial."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "id": "697568e92bd6"
+   },
+   "outputs": [],
+   "source": [
+    "from datetime import datetime\n",
+    "\n",
+    "TIMESTAMP = datetime.now().strftime(\"%Y%m%d%H%M%S\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "id": "dr--iN2kAylZ"
+   },
+   "source": [
+    "### Authenticate your Google Cloud account\n",
+    "\n",
+    "**If you are using AI Platform Notebooks**, your environment is already\n",
+    "authenticated. Skip this step."
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "id": "sBCra4QMA2wR"
+   },
+   "source": [
+    "**If you are using Colab**, run the cell below and follow the instructions\n",
+    "when prompted to authenticate your account via oAuth.\n",
+    "\n",
+    "**Otherwise**, follow these steps:\n",
+    "\n",
+    "1. In the Cloud Console, go to the [**Create service account key**\n",
+    "   page](https://console.cloud.google.com/apis/credentials/serviceaccountkey).\n",
+    "\n",
+    "2. Click **Create service account**.\n",
+    "\n",
+    "3. In the **Service account name** field, enter a name, and\n",
+    "   click **Create**.\n",
+    "\n",
+    "4. In the **Grant this service account access to project** section, click the **Role** drop-down list. Type \"AI Platform\"\n",
+    "into the filter box, and select\n",
+    "   **AI Platform Administrator**. Type \"Storage Object Admin\" into the filter box, and select **Storage Object Admin**.\n",
+    "\n",
+    "5. Click *Create*. A JSON file that contains your key downloads to your\n",
+    "local environment.\n",
+    "\n",
+    "6. Enter the path to your service account key as the\n",
+    "`GOOGLE_APPLICATION_CREDENTIALS` variable in the cell below and run the cell."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "id": "PyQmSRbKA8r-"
+   },
+   "outputs": [],
+   "source": [
+    "import os\n",
+    "import sys\n",
+    "\n",
+    "# If you are running this notebook in Colab, run this cell and follow the\n",
+    "# instructions to authenticate your GCP account. This provides access to your\n",
+    "# Cloud Storage bucket and lets you submit training jobs and prediction\n",
+    "# requests.\n",
+    "\n",
+    "# If on Google Cloud Notebooks, then don't execute this code\n",
+    "IS_GOOGLE_CLOUD_NOTEBOOK = os.path.exists(\"/opt/deeplearning/metadata/env_version\")\n",
+    "\n",
+    "if not IS_GOOGLE_CLOUD_NOTEBOOK:\n",
+    "    if \"google.colab\" in sys.modules:\n",
+    "        from google.colab import auth as google_auth\n",
+    "\n",
+    "        google_auth.authenticate_user()\n",
+    "\n",
+    "    # If you are running this notebook locally, replace the string below with the\n",
+    "    # path to your service account key and run this cell to authenticate your GCP\n",
+    "    # account.\n",
+    "    elif not os.getenv(\"IS_TESTING\"):\n",
+    "        %env GOOGLE_APPLICATION_CREDENTIALS ''"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "id": "XoEqT2Y4DJmf"
+   },
+   "source": [
+    "### Import libraries and define constants"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "id": "Y9Uo3tifg1kx"
+   },
+   "source": [
+    "Import required libraries."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "id": "pRUOFELefqf1"
+   },
+   "outputs": [],
+   "source": [
+    "import matplotlib.pyplot as plt\n",
+    "import pandas as pd\n",
+    "from google.cloud import aiplatform\n",
+    "from tensorflow.python.keras import Sequential, layers\n",
+    "from tensorflow.python.lib.io import file_io"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "id": "xtXZWmYqJ1bh"
+   },
+   "source": [
+    "Define some constants"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "id": "JIOrI-hoJ46P"
+   },
+   "outputs": [],
+   "source": [
+    "EXPERIMENT_NAME = \"\"  # @param {type:\"string\"}\n",
+    "REGION = \"[your-region]\"  # @param {type:\"string\"}"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "id": "jWQLXXNVN4Lv"
+   },
+   "source": [
+    "If EXEPERIMENT_NAME is not set, set a default one below:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "id": "Q1QInYWOKsmo"
+   },
+   "outputs": [],
+   "source": [
+    "if EXPERIMENT_NAME == \"\" or EXPERIMENT_NAME is None:\n",
+    "    EXPERIMENT_NAME = \"my-experiment-\" + TIMESTAMP"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "id": "Xuny18aMcWDb"
+   },
+   "source": [
+    "## Concepts\n",
+    "\n",
+    "To better understanding how parameters and metrics are stored and organized, we'd like to introduce the following concepts:\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "id": "NThDci5bp0Uw"
+   },
+   "source": [
+    "### Experiment\n",
+    "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."
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "id": "SAyRR3Ydp4X5"
+   },
+   "source": [
+    "### Run\n",
+    "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. "
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "id": "l1YW2pgyegFP"
+   },
+   "source": [
+    "## Getting started tracking parameters and metrics\n",
+    "\n",
+    "You can use the Vertex AI SDK to track metrics and parameters for models trained locally. \n",
+    "\n",
+    "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](https://www.kaggle.com/devanshbesain/exploration-and-analysis-auto-mpg)."
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "id": "KPY41M9_AhZU"
+   },
+   "source": [
+    "### Load and process the training dataset"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "id": "bfMQSmRuUuX-"
+   },
+   "source": [
+    "Download and process the dataset."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "id": "RiQuMv4bmpuV"
+   },
+   "outputs": [],
+   "source": [
+    "def read_data(file_path):\n",
+    "    column_names = [\n",
+    "        \"MPG\",\n",
+    "        \"Cylinders\",\n",
+    "        \"Displacement\",\n",
+    "        \"Horsepower\",\n",
+    "        \"Weight\",\n",
+    "        \"Acceleration\",\n",
+    "        \"Model Year\",\n",
+    "        \"Origin\",\n",
+    "    ]\n",
+    "    with file_io.FileIO(file_path, \"r\") as f:\n",
+    "        raw_dataset = pd.read_csv(\n",
+    "            f,\n",
+    "            names=column_names,\n",
+    "            na_values=\"?\",\n",
+    "            comment=\"\\t\",\n",
+    "            sep=\" \",\n",
+    "            skipinitialspace=True,\n",
+    "        )\n",
+    "    dataset = raw_dataset.dropna()\n",
+    "    dataset[\"Origin\"] = dataset[\"Origin\"].map(\n",
+    "        lambda x: {1: \"USA\", 2: \"Europe\", 3: \"Japan\"}.get(x)\n",
+    "    )\n",
+    "    dataset = pd.get_dummies(dataset, prefix=\"\", prefix_sep=\"\")\n",
+    "    return dataset\n",
+    "\n",
+    "\n",
+    "dataset = read_data(\"gs://cloud-samples-data/ai-platform/auto_mpg/auto-mpg.data\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "id": "Y06J7A7yU21t"
+   },
+   "source": [
+    "Split dataset for training and testing."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "id": "p5JBCBKyH-NC"
+   },
+   "outputs": [],
+   "source": [
+    "def train_test_split(dataset, split_frac=0.8, random_state=0):\n",
+    "    train_dataset = dataset.sample(frac=split_frac, random_state=random_state)\n",
+    "    test_dataset = dataset.drop(train_dataset.index)\n",
+    "    train_labels = train_dataset.pop(\"MPG\")\n",
+    "    test_labels = test_dataset.pop(\"MPG\")\n",
+    "\n",
+    "    return train_dataset, test_dataset, train_labels, test_labels\n",
+    "\n",
+    "\n",
+    "train_dataset, test_dataset, train_labels, test_labels = train_test_split(dataset)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "id": "gaNNTFPaU7KT"
+   },
+   "source": [
+    "Normalize the features in the dataset for better model performance."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "id": "VGq5QCoyIEWJ"
+   },
+   "outputs": [],
+   "source": [
+    "def normalize_dataset(train_dataset, test_dataset):\n",
+    "    train_stats = train_dataset.describe()\n",
+    "    train_stats = train_stats.transpose()\n",
+    "\n",
+    "    def norm(x):\n",
+    "        return (x - train_stats[\"mean\"]) / train_stats[\"std\"]\n",
+    "\n",
+    "    normed_train_data = norm(train_dataset)\n",
+    "    normed_test_data = norm(test_dataset)\n",
+    "\n",
+    "    return normed_train_data, normed_test_data\n",
+    "\n",
+    "\n",
+    "normed_train_data, normed_test_data = normalize_dataset(train_dataset, test_dataset)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "id": "UBXUgxgqA_GB"
+   },
+   "source": [
+    "### Define ML model and training function"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "id": "66odBYKrIN4q"
+   },
+   "outputs": [],
+   "source": [
+    "def train(\n",
+    "    train_data,\n",
+    "    train_labels,\n",
+    "    num_units=64,\n",
+    "    activation=\"relu\",\n",
+    "    dropout_rate=0.0,\n",
+    "    validation_split=0.2,\n",
+    "    epochs=1000,\n",
+    "):\n",
+    "\n",
+    "    model = Sequential(\n",
+    "        [\n",
+    "            layers.Dense(\n",
+    "                num_units,\n",
+    "                activation=activation,\n",
+    "                input_shape=[len(train_dataset.keys())],\n",
+    "            ),\n",
+    "            layers.Dropout(rate=dropout_rate),\n",
+    "            layers.Dense(num_units, activation=activation),\n",
+    "            layers.Dense(1),\n",
+    "        ]\n",
+    "    )\n",
+    "\n",
+    "    model.compile(loss=\"mse\", optimizer=\"adam\", metrics=[\"mae\", \"mse\"])\n",
+    "    print(model.summary())\n",
+    "\n",
+    "    history = model.fit(\n",
+    "        train_data, train_labels, epochs=epochs, validation_split=validation_split\n",
+    "    )\n",
+    "\n",
+    "    return model, history"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "id": "O8XJZB3gR8eL"
+   },
+   "source": [
+    "### Initialize the Model Builder SDK and create an Experiment\n",
+    "\n",
+    "Initialize the *client* for Vertex AI and create an experiment."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "id": "o_wnT10RJ7-W"
+   },
+   "outputs": [],
+   "source": [
+    "aiplatform.init(project=PROJECT_ID, location=REGION, experiment=EXPERIMENT_NAME)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "id": "u-iTnzt3B6Z_"
+   },
+   "source": [
+    "### Start several model training runs\n",
+    "\n",
+    "Training parameters and metrics are logged for each run."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "id": "i2wnpu8_7JfV"
+   },
+   "outputs": [],
+   "source": [
+    "parameters = [\n",
+    "    {\"num_units\": 16, \"epochs\": 3, \"dropout_rate\": 0.1},\n",
+    "    {\"num_units\": 16, \"epochs\": 10, \"dropout_rate\": 0.1},\n",
+    "    {\"num_units\": 16, \"epochs\": 10, \"dropout_rate\": 0.2},\n",
+    "    {\"num_units\": 32, \"epochs\": 10, \"dropout_rate\": 0.1},\n",
+    "    {\"num_units\": 32, \"epochs\": 10, \"dropout_rate\": 0.2},\n",
+    "]\n",
+    "\n",
+    "for i, params in enumerate(parameters):\n",
+    "    aiplatform.start_run(run=f\"auto-mpg-local-run-{i}\")\n",
+    "    aiplatform.log_params(params)\n",
+    "    model, history = train(\n",
+    "        normed_train_data,\n",
+    "        train_labels,\n",
+    "        num_units=params[\"num_units\"],\n",
+    "        activation=\"relu\",\n",
+    "        epochs=params[\"epochs\"],\n",
+    "        dropout_rate=params[\"dropout_rate\"],\n",
+    "    )\n",
+    "    aiplatform.log_metrics(\n",
+    "        {metric: values[-1] for metric, values in history.history.items()}\n",
+    "    )\n",
+    "\n",
+    "    loss, mae, mse = model.evaluate(normed_test_data, test_labels, verbose=2)\n",
+    "    aiplatform.log_metrics({\"eval_loss\": loss, \"eval_mae\": mae, \"eval_mse\": mse})"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "id": "jZLrJZTfL7tE"
+   },
+   "source": [
+    "### Extract parameters and metrics into a dataframe for analysis"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "id": "A1PqKxlpOZa2"
+   },
+   "source": [
+    "We can also extract all parameters and metrics associated with any Experiment into a dataframe for further analysis."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "id": "jbRf1WoH_vbY"
+   },
+   "outputs": [],
+   "source": [
+    "experiment_df = aiplatform.get_experiment_df()\n",
+    "experiment_df"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "id": "EYuYgqVCMKU1"
+   },
+   "source": [
+    "### Visualizing an experiment's parameters and metrics"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "id": "r8orCj8iJuO1"
+   },
+   "outputs": [],
+   "source": [
+    "plt.rcParams[\"figure.figsize\"] = [15, 5]\n",
+    "\n",
+    "ax = pd.plotting.parallel_coordinates(\n",
+    "    experiment_df.reset_index(level=0),\n",
+    "    \"run_name\",\n",
+    "    cols=[\n",
+    "        \"param.num_units\",\n",
+    "        \"param.dropout_rate\",\n",
+    "        \"param.epochs\",\n",
+    "        \"metric.loss\",\n",
+    "        \"metric.val_loss\",\n",
+    "        \"metric.eval_loss\",\n",
+    "    ],\n",
+    "    color=[\"blue\", \"green\", \"pink\", \"red\"],\n",
+    ")\n",
+    "ax.set_yscale(\"symlog\")\n",
+    "ax.legend(bbox_to_anchor=(1.0, 0.5))"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "id": "WTHvPMweMlP1"
+   },
+   "source": [
+    "## Visualizing experiments in Cloud Console"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "id": "F19_5lw0MqXv"
+   },
+   "source": [
+    "Run the following to get the URL of Vertex AI Experiments for your project.\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "id": "GmN9vE9pqqzt"
+   },
+   "outputs": [],
+   "source": [
+    "print(\"Vertex AI Experiments:\")\n",
+    "print(\n",
+    "    f\"https://console.cloud.google.com/ai/platform/experiments/experiments?folder=&organizationId=&project={PROJECT_ID}\"\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "id": "TpV-iwP9qw9c"
+   },
+   "source": [
+    "## Cleaning up\n",
+    "\n",
+    "To clean up all Google Cloud resources used in this project, you can [delete the Google Cloud\n",
+    "project](https://cloud.google.com/resource-manager/docs/creating-managing-projects#shutting_down_projects) you used for the tutorial."
+   ]
+  }
+ ],
+ "metadata": {
+  "colab": {
+   "collapsed_sections": [],
+   "name": "sdk-metric-parameter-tracking-for-locally-trained-models.ipynb",
+   "toc_visible": true
+  },
+  "kernelspec": {
+   "display_name": "Python 3",
+   "name": "python3"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 0
+}