Merge branch 'main' into autlogging

notebooks/community/ml_ops/stage2/get_started_with_vertex_experiments_autologging_tf.ipynb

@@ -0,0 +1,974 @@
+{
+  "cells": [
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "copyright"
+      },
+      "outputs": [],
+      "source": [
+        "# Copyright 2022 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": "title"
+      },
+      "source": [
+        "# E2E ML on GCP: MLOps stage 2 : Get started with autologging using Vertex AI Experiments for TensorFlow models\n",
+        "\n",
+        "<table align=\"left\">\n",
+        "  <td>\n",
+        "    <a href=\"https://colab.research.google.com/github/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/community/ml_ops/stage6/get_started_with_experiments_autologging_tf.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/main/notebooks/community/ml_ops/stage6/get_started_with_experiments_autologging_tf.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",
+        "  <td>\n",
+        "    <a href=\"https://console.cloud.google.com/vertex-ai/workbench/deploy-notebook?download_url=https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/main/notebooks/community/ml_ops/stage6/get_started_with_experiments_autologging_tf.ipynb\">\n",
+        "      <img src=\"https://lh3.googleusercontent.com/UiNooY4LUgW_oTvpsNhPpQzsstV5W8F7rYgxgGBD85cWJoLmrOzhVs_ksK_vgx40SHs7jCqkTkCk=e14-rj-sc0xffffff-h130-w32\" alt=\"Vertex AI logo\">\n",
+        "      Open in Vertex AI Workbench\n",
+        "    </a>\n",
+        "  </td>\n",
+        "</table>\n",
+        "<br/><br/><br/>"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "overview:automl"
+      },
+      "source": [
+        "## Overview\n",
+        "\n",
+        "\n",
+        "This tutorial demonstrates how to use the `Vertex AI Experiments` with DIY code to implement automatic logging of parameters and metrics for experiments."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "objective:automl,training,batch_prediction"
+      },
+      "source": [
+        "### Objective\n",
+        "\n",
+        "In this tutorial, you learn how to create an experiment for training a TensorFlow model, and automatically log parameters and metrics using the enclosed do-it-yourself (DIY) code.\n",
+        "\n",
+        "This tutorial uses the following Google Cloud ML services and resources:\n",
+        "\n",
+        "- `Vertex AI Experiments`\n",
+        "\n",
+        "The steps performed include:\n",
+        "\n",
+        "- Construct the DIY autologging code.\n",
+        "- Construct training package with call to autologging.\n",
+        "- Train a model.\n",
+        "- View the experiment\n",
+        "- Delete the experiment."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "2739272aae1b"
+      },
+      "source": [
+        "### Model\n",
+        "\n",
+        "The model used for this tutorial is a pretrain TensorFlow model that was trained on the [Boston Housing Prices dataset](https://www.cs.toronto.edu/~delve/data/boston/bostonDetail.html). The version of the dataset this tutorial is built into TensorFlow. The trained model predicts the median price of a house in units of 1K USD."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "costs"
+      },
+      "source": [
+        "### Costs\n",
+        "\n",
+        "This tutorial uses billable components of Google Cloud:\n",
+        "\n",
+        "* Vertex AI\n",
+        "* Cloud Storage\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": "setup_local"
+      },
+      "source": [
+        "### Set up your local development environment\n",
+        "\n",
+        "If you are using Colab or Vertex Workbench AI Notebooks, your environment already meets all the requirements to run this notebook. You can skip this step.\n",
+        "\n",
+        "Otherwise, make sure your environment meets this notebook's requirements. You need the following:\n",
+        "\n",
+        "- The Cloud Storage SDK\n",
+        "- Git\n",
+        "- Python 3\n",
+        "- virtualenv\n",
+        "- Jupyter notebook running in a virtual environment with Python 3\n",
+        "\n",
+        "The Cloud Storage guide to [Setting up a Python development environment](https://cloud.google.com/python/setup) and the [Jupyter installation guide](https://jupyter.org/install) provide detailed instructions for meeting these requirements. The following steps provide a condensed set of instructions:\n",
+        "\n",
+        "1. [Install and initialize the SDK](https://cloud.google.com/sdk/docs/).\n",
+        "\n",
+        "2. [Install Python 3](https://cloud.google.com/python/setup#installing_python).\n",
+        "\n",
+        "3. [Install virtualenv](https://cloud.google.com/python/setup#installing_and_using_virtualenv) and create a virtual environment that uses Python 3.  Activate the virtual environment.\n",
+        "\n",
+        "4. To install Jupyter, run `pip3 install jupyter` on the command-line in a terminal shell.\n",
+        "\n",
+        "5. To launch Jupyter, run `jupyter notebook` on the command-line in a terminal shell.\n",
+        "\n",
+        "6. Open this notebook in the Jupyter Notebook Dashboard.\n"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "install_aip:mbsdk"
+      },
+      "source": [
+        "## Installation\n",
+        "\n",
+        "Install the following packages to execute this notebook."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "install_aip:mbsdk"
+      },
+      "outputs": [],
+      "source": [
+        "import os\n",
+        "\n",
+        "# The Vertex AI Workbench Notebook product has specific requirements\n",
+        "IS_WORKBENCH_NOTEBOOK = os.getenv(\"DL_ANACONDA_HOME\")\n",
+        "IS_USER_MANAGED_WORKBENCH_NOTEBOOK = os.path.exists(\n",
+        "    \"/opt/deeplearning/metadata/env_version\"\n",
+        ")\n",
+        "\n",
+        "# Vertex AI Notebook requires dependencies to be installed with '--user'\n",
+        "USER_FLAG = \"\"\n",
+        "if IS_WORKBENCH_NOTEBOOK:\n",
+        "    USER_FLAG = \"--user\"\n",
+        "\n",
+        "! pip3 install --upgrade --quiet {USER_FLAG} google-cloud-aiplatform \\\n",
+        "                                             tensorflow==2.5 \\\n",
+        "                                             numpy"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "restart"
+      },
+      "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": "restart"
+      },
+      "outputs": [],
+      "source": [
+        "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": "before_you_begin:nogpu"
+      },
+      "source": [
+        "## Before you begin\n",
+        "\n",
+        "### GPU runtime\n",
+        "\n",
+        "This tutorial does not require a GPU runtime.\n",
+        "\n",
+        "### 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",
+        "2. [Make sure that billing is enabled for your project.](https://cloud.google.com/billing/docs/how-to/modify-project)\n",
+        "\n",
+        "3. [Enable the following APIs: Vertex AI APIs, Compute Engine APIs, and Cloud Storage.](https://console.cloud.google.com/flows/enableapi?apiid=ml.googleapis.com,compute_component,storage-component.googleapis.com)\n",
+        "\n",
+        "4. If you are running this notebook locally, you need to install the [Cloud SDK]((https://cloud.google.com/sdk)).\n",
+        "\n",
+        "5. 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 `$`."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "project_id"
+      },
+      "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": "set_project_id"
+      },
+      "outputs": [],
+      "source": [
+        "PROJECT_ID = \"[your-project-id]\"  # @param {type:\"string\"}"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "autoset_project_id"
+      },
+      "outputs": [],
+      "source": [
+        "if PROJECT_ID == \"\" or PROJECT_ID is None or PROJECT_ID == \"[your-project-id]\":\n",
+        "    # Get your GCP project id from gcloud\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": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "set_gcloud_project_id"
+      },
+      "outputs": [],
+      "source": [
+        "! gcloud config set project $PROJECT_ID"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "region"
+      },
+      "source": [
+        "#### Region\n",
+        "\n",
+        "You can also change the `REGION` variable, which is used for operations\n",
+        "throughout the rest of this notebook.  Below are regions supported for Vertex AI. We recommend that you choose the region closest to you.\n",
+        "\n",
+        "- Americas: `us-central1`\n",
+        "- Europe: `europe-west4`\n",
+        "- Asia Pacific: `asia-east1`\n",
+        "\n",
+        "You may not use a multi-regional bucket for training with Vertex AI. Not all regions provide support for all Vertex AI services.\n",
+        "\n",
+        "Learn more about [Vertex AI regions](https://cloud.google.com/vertex-ai/docs/general/locations)"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "region"
+      },
+      "outputs": [],
+      "source": [
+        "REGION = \"[your-region]\"  # @param {type: \"string\"}\n",
+        "\n",
+        "if REGION == \"[your-region]\":\n",
+        "    REGION = \"us-central1\""
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "timestamp"
+      },
+      "source": [
+        "#### UUID\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 uuid 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": "timestamp"
+      },
+      "outputs": [],
+      "source": [
+        "import random\n",
+        "import string\n",
+        "\n",
+        "\n",
+        "# Generate a uuid of a specifed length(default=8)\n",
+        "def generate_uuid(length: int = 8) -> str:\n",
+        "    return \"\".join(random.choices(string.ascii_lowercase + string.digits, k=length))\n",
+        "\n",
+        "\n",
+        "UUID = generate_uuid()"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "gcp_authenticate"
+      },
+      "source": [
+        "### Authenticate your Google Cloud account\n",
+        "\n",
+        "**If you are using Vertex AI Workbench Notebooks**, your environment is already authenticated. \n",
+        "\n",
+        "**If you are using Colab**, run the cell below and follow the instructions when prompted to authenticate your account via oAuth.\n",
+        "\n",
+        "**Otherwise**, follow these steps:\n",
+        "\n",
+        "In the Cloud Console, go to the [Create service account key](https://console.cloud.google.com/apis/credentials/serviceaccountkey) page.\n",
+        "\n",
+        "**Click Create service account**.\n",
+        "\n",
+        "In the **Service account name** field, enter a name, and click **Create**.\n",
+        "\n",
+        "In the **Grant this service account access to project** section, click the Role drop-down list. Type \"Vertex\" into the filter box, and select **Vertex Administrator**. Type \"Storage Object Admin\" into the filter box, and select **Storage Object Admin**.\n",
+        "\n",
+        "Click Create. A JSON file that contains your key downloads to your local environment.\n",
+        "\n",
+        "Enter the path to your service account key as the GOOGLE_APPLICATION_CREDENTIALS variable in the cell below and run the cell."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "gcp_authenticate"
+      },
+      "outputs": [],
+      "source": [
+        "# 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",
+        "import os\n",
+        "import sys\n",
+        "\n",
+        "# If on Vertex AI Workbench, then don't execute this code\n",
+        "IS_COLAB = \"google.colab\" in sys.modules\n",
+        "if not os.path.exists(\"/opt/deeplearning/metadata/env_version\") and not os.getenv(\n",
+        "    \"DL_ANACONDA_HOME\"\n",
+        "):\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": "setup_vars"
+      },
+      "source": [
+        "### Import libraries and define constants"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "59963fb7178f"
+      },
+      "outputs": [],
+      "source": [
+        "import google.cloud.aiplatform as aiplatform\n",
+        "import numpy as np\n",
+        "import tensorflow as tf"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "init_aip:mbsdk"
+      },
+      "source": [
+        "## Initialize Vertex AI SDK for Python\n",
+        "\n",
+        "Initialize the Vertex AI SDK for Python for your project and corresponding bucket."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "init_aip:mbsdk"
+      },
+      "outputs": [],
+      "source": [
+        "aiplatform.init(project=PROJECT_ID, location=REGION)"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "ae8f31c8c617"
+      },
+      "source": [
+        "## DIY code for autologging TensorFlow Keras models\n",
+        "\n",
+        "The code below implements autologging for TensorFlow models using the Keras API.\n",
+        "\n",
+        "- `autologging()`: Initializes the experiment and uses heap injection to replace TF.keras `Sequential` and `Model` symbols on the heap with the redirect wrapper classes `VertexTFSequential` and `VertexTFModel`, respectively.\n",
+        "\n",
+        "- `VertexTFSequential`: A subclass of the tf.keras.Sequential class.\n",
+        "    - `compile()`: overridden method of super class. Automatically logs specified hyperparameters and calls the underlying `compile()` method.\n",
+        "    - `fit()`: overridden method of super class. Automatically logs specified hyperparameters, calls the underlying `fit()` method, and logs the resulting metrics.\n",
+        "    - `evaluate()`: overridden method of super class. Calls the underlying `evaluate()` method, and logs the resulting metrics.\n",
+        "- `VertexTFModel`: A subclass of the tf.keras.Model class."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "8eb012e5d7ef"
+      },
+      "outputs": [],
+      "source": [
+        "def autolog(\n",
+        "    project: str = None,\n",
+        "    location: str = None,\n",
+        "    staging_bucket: str = None,\n",
+        "    experiment: str = None,\n",
+        "    run: str = None,\n",
+        "    framework: str = \"tf\",\n",
+        "):\n",
+        "    \"\"\"\n",
+        "    Enable automatic logging of parameters and metrics in Vertex AI Experiments,\n",
+        "    for corresponding framework.\n",
+        "\n",
+        "        project: The project ID\n",
+        "        location : The region\n",
+        "        staging_bucket: temporary bucket\n",
+        "        experiment: The name of the experiment\n",
+        "        run: The name of the run within the experiment\n",
+        "        framework: The ML framework for which a model is being trained.\n",
+        "    \"\"\"\n",
+        "    # autologging\n",
+        "    if framework == \"tf\":\n",
+        "        try:\n",
+        "            globals()[\"Sequential\"] = VertexTFSequential\n",
+        "            if \"tf\" in globals():\n",
+        "                tf.keras.Sequential = VertexTFSequential\n",
+        "            if \"tensorflow\" in globals():\n",
+        "                tensorflow.keras.Sequential = VertexTFSequential\n",
+        "        except:\n",
+        "            pass\n",
+        "\n",
+        "        try:\n",
+        "            globals()[\"Model\"] = VertexTFModel\n",
+        "            if \"tf\" in globals():\n",
+        "                tf.keras.Model = VertexTFModel\n",
+        "            if \"tensorflow\" in globals():\n",
+        "                tensorflow.keras.Model = VertexTFModel\n",
+        "        except:\n",
+        "            pass\n",
+        "\n",
+        "    if project:\n",
+        "        aiplatform.init(\n",
+        "            project=project, location=location, staging_bucket=staging_bucket\n",
+        "        )\n",
+        "\n",
+        "    if experiment:\n",
+        "        aiplatform.init(experiment=experiment)\n",
+        "    if run:\n",
+        "        aiplatform.start_run(run)\n",
+        "\n",
+        "\n",
+        "class VertexTFSequential(tf.keras.Sequential):\n",
+        "    \"\"\"\n",
+        "    Sublass of the tensorflow.keras.Sequential model type. Overrides with pass thru of\n",
+        "    underlying super class methods to automatically log parameters/metrics for Vertex AI experiments.\n",
+        "\n",
+        "        compile():\n",
+        "        fit():\n",
+        "        evaluate():\n",
+        "    \"\"\"\n",
+        "\n",
+        "    def __init__(self, layers):\n",
+        "        return super().__init__(layers)\n",
+        "\n",
+        "    def compile(\n",
+        "        self,\n",
+        "        optimizer=\"rmsprop\",\n",
+        "        loss=None,\n",
+        "        metrics=None,\n",
+        "        loss_weights=None,\n",
+        "        weighted_metrics=None,\n",
+        "        run_eagerly=None,\n",
+        "        steps_per_execution=None,\n",
+        "    ):\n",
+        "        try:\n",
+        "            learning_rate = optimizer.learning_rate.numpy()\n",
+        "            aiplatform.log_params({\"train.learning_rate\": float(learning_rate)})\n",
+        "        except:\n",
+        "            pass\n",
+        "        return super().compile(\n",
+        "            loss=loss,\n",
+        "            optimizer=optimizer,\n",
+        "            metrics=metrics,\n",
+        "            loss_weights=loss_weights,\n",
+        "            weighted_metrics=weighted_metrics,\n",
+        "            run_eagerly=run_eagerly,\n",
+        "            steps_per_execution=steps_per_execution,\n",
+        "        )\n",
+        "\n",
+        "    def fit(\n",
+        "        self,\n",
+        "        x=None,\n",
+        "        y=None,\n",
+        "        batch_size=None,\n",
+        "        epochs=1,\n",
+        "        verbose=\"auto\",\n",
+        "        callbacks=None,\n",
+        "        validation_split=0.0,\n",
+        "        validation_data=None,\n",
+        "        shuffle=True,\n",
+        "        class_weight=None,\n",
+        "        sample_weight=None,\n",
+        "        initial_epoch=0,\n",
+        "        steps_per_epoch=None,\n",
+        "        validation_steps=None,\n",
+        "        validation_batch_size=None,\n",
+        "        validation_freq=1,\n",
+        "        max_queue_size=10,\n",
+        "        workers=1,\n",
+        "        use_multiprocessing=False,\n",
+        "    ):\n",
+        "        aiplatform.log_params({\"train.epochs\": int(epochs)})\n",
+        "        if batch_size:\n",
+        "            aiplatform.log_params({\"train.batch_size\": int(batch_size)})\n",
+        "        if steps_per_epoch:\n",
+        "            aiplatform.log_params({\"train.steps\": int(steps_per_epoch)})\n",
+        "\n",
+        "        history = super().fit(\n",
+        "            x=x,\n",
+        "            y=y,\n",
+        "            batch_size=batch_size,\n",
+        "            epochs=epochs,\n",
+        "            verbose=verbose,\n",
+        "            callbacks=callbacks,\n",
+        "            validation_split=validation_split,\n",
+        "            validation_data=validation_data,\n",
+        "            shuffle=shuffle,\n",
+        "            class_weight=class_weight,\n",
+        "            sample_weight=sample_weight,\n",
+        "            initial_epoch=initial_epoch,\n",
+        "            steps_per_epoch=steps_per_epoch,\n",
+        "            validation_steps=validation_steps,\n",
+        "            validation_batch_size=validation_batch_size,\n",
+        "            validation_freq=validation_freq,\n",
+        "            max_queue_size=max_queue_size,\n",
+        "            workers=workers,\n",
+        "            use_multiprocessing=use_multiprocessing,\n",
+        "        )\n",
+        "\n",
+        "        for key, val in history.history.items():\n",
+        "            aiplatform.log_metrics({f\"train.{key}\": val[-1]})\n",
+        "        return history\n",
+        "\n",
+        "    def evaluate(\n",
+        "        self,\n",
+        "        x=None,\n",
+        "        y=None,\n",
+        "        batch_size=None,\n",
+        "        verbose=1,\n",
+        "        sample_weight=None,\n",
+        "        steps=None,\n",
+        "        callbacks=None,\n",
+        "        max_queue_size=10,\n",
+        "        workers=1,\n",
+        "        use_multiprocessing=False,\n",
+        "        return_dict=False,\n",
+        "    ):\n",
+        "\n",
+        "        metrics = super().evaluate(\n",
+        "            x=x,\n",
+        "            y=y,\n",
+        "            batch_size=batch_size,\n",
+        "            verbose=verbose,\n",
+        "            sample_weight=sample_weight,\n",
+        "            steps=steps,\n",
+        "            callbacks=callbacks,\n",
+        "            max_queue_size=max_queue_size,\n",
+        "            workers=workers,\n",
+        "            use_multiprocessing=use_multiprocessing,\n",
+        "            return_dict=return_dict,\n",
+        "        )\n",
+        "\n",
+        "        aiplatform.log_metrics({\"eval.loss\": metrics[0]})\n",
+        "        for _ in range(1, len(metrics)):\n",
+        "            aiplatform.log_metrics({\"eval.metric\": metrics[_]})\n",
+        "        return metrics\n",
+        "\n",
+        "\n",
+        "class VertexTFModel(tf.keras.Model):\n",
+        "    \"\"\"\n",
+        "    Sublass of the tensorflow.keras.Model model type. Overrides with pass thru of\n",
+        "    underlying super class methods to automatically log parameters/metrics for Vertex AI experiments.\n",
+        "\n",
+        "        compile():\n",
+        "        fit():\n",
+        "        evaluate():\n",
+        "    \"\"\"\n",
+        "\n",
+        "    def __init__(self, inputs, outputs):\n",
+        "        return super().__init__(inputs, outputs)\n",
+        "\n",
+        "    def compile(\n",
+        "        self,\n",
+        "        optimizer=\"rmsprop\",\n",
+        "        loss=None,\n",
+        "        metrics=None,\n",
+        "        loss_weights=None,\n",
+        "        weighted_metrics=None,\n",
+        "        run_eagerly=None,\n",
+        "        steps_per_execution=None,\n",
+        "    ):\n",
+        "        try:\n",
+        "            learning_rate = optimizer.learning_rate.numpy()\n",
+        "            aiplatform.log_params({\"train.learning_rate\": float(learning_rate)})\n",
+        "        except:\n",
+        "            pass\n",
+        "        return super().compile(\n",
+        "            loss=loss,\n",
+        "            optimizer=optimizer,\n",
+        "            metrics=metrics,\n",
+        "            loss_weights=loss_weights,\n",
+        "            weighted_metrics=weighted_metrics,\n",
+        "            run_eagerly=run_eagerly,\n",
+        "            steps_per_execution=steps_per_execution,\n",
+        "        )\n",
+        "\n",
+        "    def fit(\n",
+        "        self,\n",
+        "        x=None,\n",
+        "        y=None,\n",
+        "        batch_size=None,\n",
+        "        epochs=1,\n",
+        "        verbose=\"auto\",\n",
+        "        callbacks=None,\n",
+        "        validation_split=0.0,\n",
+        "        validation_data=None,\n",
+        "        shuffle=True,\n",
+        "        class_weight=None,\n",
+        "        sample_weight=None,\n",
+        "        initial_epoch=0,\n",
+        "        steps_per_epoch=None,\n",
+        "        validation_steps=None,\n",
+        "        validation_batch_size=None,\n",
+        "        validation_freq=1,\n",
+        "        max_queue_size=10,\n",
+        "        workers=1,\n",
+        "        use_multiprocessing=False,\n",
+        "    ):\n",
+        "        aiplatform.log_params({\"train.epochs\": int(epochs)})\n",
+        "        if batch_size:\n",
+        "            aiplatform.log_params({\"train.batch_size\": int(batch_size)})\n",
+        "        if steps_per_epoch:\n",
+        "            aiplatform.log_params({\"train.steps\": int(steps_per_epoch)})\n",
+        "\n",
+        "        history = super().fit(\n",
+        "            x=x,\n",
+        "            y=y,\n",
+        "            batch_size=batch_size,\n",
+        "            epochs=epochs,\n",
+        "            verbose=verbose,\n",
+        "            callbacks=callbacks,\n",
+        "            validation_split=validation_split,\n",
+        "            validation_data=validation_data,\n",
+        "            shuffle=shuffle,\n",
+        "            class_weight=class_weight,\n",
+        "            sample_weight=sample_weight,\n",
+        "            initial_epoch=initial_epoch,\n",
+        "            steps_per_epoch=steps_per_epoch,\n",
+        "            validation_steps=validation_steps,\n",
+        "            validation_batch_size=validation_batch_size,\n",
+        "            validation_freq=validation_freq,\n",
+        "            max_queue_size=max_queue_size,\n",
+        "            workers=workers,\n",
+        "            use_multiprocessing=use_multiprocessing,\n",
+        "        )\n",
+        "\n",
+        "        for key, val in history.history.items():\n",
+        "            aiplatform.log_metrics({f\"train.{key}\": val[-1]})\n",
+        "        return history\n",
+        "\n",
+        "    def evaluate(\n",
+        "        self,\n",
+        "        x=None,\n",
+        "        y=None,\n",
+        "        batch_size=None,\n",
+        "        verbose=1,\n",
+        "        sample_weight=None,\n",
+        "        steps=None,\n",
+        "        callbacks=None,\n",
+        "        max_queue_size=10,\n",
+        "        workers=1,\n",
+        "        use_multiprocessing=False,\n",
+        "        return_dict=False,\n",
+        "    ):\n",
+        "\n",
+        "        metrics = super().evaluate(\n",
+        "            x=x,\n",
+        "            y=y,\n",
+        "            batch_size=batch_size,\n",
+        "            verbose=verbose,\n",
+        "            sample_weight=sample_weight,\n",
+        "            steps=steps,\n",
+        "            callbacks=callbacks,\n",
+        "            max_queue_size=max_queue_size,\n",
+        "            workers=workers,\n",
+        "            use_multiprocessing=use_multiprocessing,\n",
+        "            return_dict=return_dict,\n",
+        "        )\n",
+        "\n",
+        "        aiplatform.log_metrics({\"eval.loss\": metrics[0]})\n",
+        "        for _ in range(1, len(metrics)):\n",
+        "            aiplatform.log_metrics({\"eval.metric\": metrics[_]})\n",
+        "        return metrics"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "ce76826902c0"
+      },
+      "source": [
+        "### Train the model with Vertex AI Experiments\n",
+        "\n",
+        "In the following code, you build, train and evaluate a TensorFlow tabular model. The Python script includes the following calls to integrate `Vertex AI Experiments`:\n",
+        "\n",
+        "- command-line arguments: The arguments `experiment` and `run` are used to pass in the experiment and run names for the experiment.\n",
+        "- `autologging()`: Initializes the experiment and does the heap injection.\n",
+        "- `aiplatform.start_execution()`: Initializes a context for linking artifacts.\n",
+        "- `aiplatform.end_run()`: Ends the experiment.\n",
+        "\n",
+        "*Note:* The initializer `Sequential` will be redirected to `VertexTFSequential` by heap injection. When subsequent calls are made to the compile(), fit() and evaluate() methods, they will be executed as the corresponding `VertexTFSequential` methods."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "427846783ed6"
+      },
+      "outputs": [],
+      "source": [
+        "EXPERIMENT_NAME = f\"myexperiment{UUID}\"\n",
+        "RUN_NAME = \"run-1\"\n",
+        "\n",
+        "\n",
+        "def make_dataset():\n",
+        "\n",
+        "    # Scaling Boston Housing data features\n",
+        "    def scale(feature):\n",
+        "        max = np.max(feature)\n",
+        "        feature = (feature / max).astype(np.float)\n",
+        "        return feature, max\n",
+        "\n",
+        "    (x_train, y_train), (x_test, y_test) = tf.keras.datasets.boston_housing.load_data(\n",
+        "        path=\"boston_housing.npz\", test_split=0.2, seed=113\n",
+        "    )\n",
+        "    params = []\n",
+        "\n",
+        "    for _ in range(13):\n",
+        "        x_train[_], max = scale(x_train[_])\n",
+        "        x_test[_], _ = scale(x_test[_])\n",
+        "        params.append(max)\n",
+        "\n",
+        "    return (x_train, y_train), (x_test, y_test)\n",
+        "\n",
+        "\n",
+        "# Build the Keras model\n",
+        "def build_and_compile_dnn_model(lr):\n",
+        "    model = tf.keras.Sequential(\n",
+        "        [\n",
+        "            tf.keras.layers.Dense(128, activation=\"relu\", input_shape=(13,)),\n",
+        "            tf.keras.layers.Dense(128, activation=\"relu\"),\n",
+        "            tf.keras.layers.Dense(1, activation=\"linear\"),\n",
+        "        ]\n",
+        "    )\n",
+        "\n",
+        "    model.compile(\n",
+        "        loss=\"mse\",\n",
+        "        optimizer=tf.keras.optimizers.RMSprop(learning_rate=lr),\n",
+        "        metrics=[tf.keras.metrics.RootMeanSquaredError()],\n",
+        "    )\n",
+        "    return model\n",
+        "\n",
+        "\n",
+        "# autologging\n",
+        "autolog(experiment=EXPERIMENT_NAME, run=RUN_NAME)\n",
+        "\n",
+        "with aiplatform.start_execution(\n",
+        "    schema_title=\"system.ContainerExecution\", display_name=\"example_training\"\n",
+        ") as execution:\n",
+        "    BATCH_SIZE = 16\n",
+        "\n",
+        "    model = build_and_compile_dnn_model(lr=0.01)\n",
+        "\n",
+        "    # Train the model\n",
+        "    (x_train, y_train), (x_test, y_test) = make_dataset()\n",
+        "    model.fit(x_train, y_train, epochs=10, batch_size=BATCH_SIZE)\n",
+        "\n",
+        "    model.evaluate(x_test, y_test)\n",
+        "\n",
+        "aiplatform.end_run()"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "5f40912e6500"
+      },
+      "source": [
+        "#### Get the experiment results\n",
+        "\n",
+        "Next, you use the experiment name as a parameter to the method `get_experiment_df()` to get the results of the experiment as a pandas dataframe."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "7e9671712230"
+      },
+      "outputs": [],
+      "source": [
+        "experiment_df = aiplatform.get_experiment_df()\n",
+        "experiment_df = experiment_df[experiment_df.experiment_name == EXPERIMENT_NAME]\n",
+        "experiment_df.T"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "e508c159d712"
+      },
+      "source": [
+        "#### Delete the experiment\n",
+        "\n",
+        "Since the experiment was created within a training script, to delete the experiment you use the `list()` method to obtain all the experiments for the project, and then filter on the experiment name."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "1a1b5fcbfde0"
+      },
+      "outputs": [],
+      "source": [
+        "experiments = aiplatform.Experiment.list()\n",
+        "for experiment in experiments:\n",
+        "    if experiment.name == EXPERIMENT_NAME:\n",
+        "        experiment.delete()"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "cleanup:mbsdk"
+      },
+      "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.\n",
+        "\n",
+        "Otherwise, you can delete the individual resources you created in this tutorial."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "9eb897e0e728"
+      },
+      "outputs": [],
+      "source": [
+        "# There are no resources to cleanup"
+      ]
+    }
+  ],
+  "metadata": {
+    "colab": {
+      "name": "get_started_with_vertex_experiments_autologging_tf.ipynb",
+      "toc_visible": true
+    },
+    "kernelspec": {
+      "display_name": "Python 3",
+      "name": "python3"
+    }
+  },
+  "nbformat": 4,
+  "nbformat_minor": 0
+}