notebooks/official/tensorboard/tensorboard_custom_training_with_prebuilt_container.ipynb

{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "JICNqM6hx1XC"
      },
      "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": "eBEO2w9My9py"
      },
      "source": [
        "# Vertex AI TensorBoard custom training with prebuilt container\n",
        "\n",
        "<table align=\"left\">\n",
        "  <td>\n",
        "    <a href=\"https://colab.research.google.com/github/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/official/tensorboard/tensorboard_custom_training_with_prebuilt_container.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/official/tensorboard/tensorboard_custom_training_with_prebuilt_container.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/official/tensorboard/tensorboard_custom_training_with_prebuilt_container.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/>\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "VL7XCFV7yCBU"
      },
      "source": [
        "## Overview\n",
        "\n",
        "### What is Vertex AI TensorBoard\n",
        "\n",
        "Vertex AI TensorBoard is an enterprise-ready managed\n",
        "version of [Open source TensorBoard](https://www.tensorflow.org/tensorboard/get_started)\n",
        "(TB), which is a Google open source project for machine learning experiment\n",
        "visualization. \n",
        "\n",
        "Vertex AI TensorBoard provides various detailed visualizations, that\n",
        "includes:\n",
        "\n",
        "*   tracking and visualizing metrics such as loss and accuracy over time,\n",
        "*   visualizing model computational graphs (ops and layers),\n",
        "*   viewing histograms of weights, biases, or other tensors as they change over time,\n",
        "*   projecting embeddings to a lower dimensional space,\n",
        "*   displaying image, text, and audio samples.\n",
        "\n",
        "In addition to the powerful visualizations from\n",
        "TensorBoard, Vertex AI TensorBoard provides:\n",
        "\n",
        "*  a persistent, shareable link to your experiment's dashboard,\n",
        "\n",
        "*  a searchable list of all experiments in a project,\n",
        "\n",
        "*  tight integrations with Vertex AI services for model training,\n",
        "\n",
        "*  enterprise-grade security, privacy, and compliance.\n",
        "\n",
        "With Vertex AI TensorBoard, you can track, visualize, and compare\n",
        "ML experiments and share them with your team.\n",
        "\n",
        "Learn more about [Vertex AI TensorBoard](https://cloud.google.com/vertex-ai/docs/experiments/tensorboard-introduction) and [Custom training](https://cloud.google.com/vertex-ai/docs/training/overview)."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "UE8vLw7SlpwE"
      },
      "source": [
        "### Objective\n",
        "\n",
        "In this tutorial, you learn how to create a custom training job using prebuilt containers, and monitor your training process on Vertex AI TensorBoard in near real time.\n",
        "\n",
        "This tutorial uses the following Google Cloud ML services and resources:\n",
        "\n",
        "- Vertex AI Training\n",
        "- Vertex AI TensorBoard\n",
        "\n",
        "The steps performed include:\n",
        "\n",
        "* Setup service account and Google Cloud Storage buckets.\n",
        "* Write your customized training code.\n",
        "* Package and upload your training code to Google Cloud Storage.\n",
        "* Create & launch your custom training job with Tensorboard enabled for near real time monitorning."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "fbF2lF8rlp3I"
      },
      "source": [
        "## Dataset\n",
        "\n",
        "The dataset used in this tutorial is the [flower dataset](https://www.tensorflow.org/datasets/catalog/tf_flowers) provided by TensorFlow. No other datasets required.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "yajq2millpnu"
      },
      "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 pricing](https://cloud.google.com/vertex-ai/pricing) and [Cloud Storage pricing](https://cloud.google.com/storage/pricing), and use the [Pricing Calculator](https://cloud.google.com/products/calculator/) to generate a cost estimate based on your projected usage."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "s3moH5AexXpk"
      },
      "source": [
        "### Install additional packages\n",
        "\n",
        "Install the following packages required to execute this notebook."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "pJwxNvcA8V_c"
      },
      "outputs": [],
      "source": [
        "! pip3 install --upgrade --quiet google-cloud-aiplatform "
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "restart"
      },
      "source": [
        "### Colab only: Uncomment the following cell to restart the kernel.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "bzPxhxS5lugp"
      },
      "outputs": [],
      "source": [
        "# Automatically restart kernel after installs so that your environment can access the new packages\n",
        "# import IPython\n",
        "\n",
        "# app = IPython.Application.instance()\n",
        "# app.kernel.do_shutdown(True)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "d2qpIurSjmpT"
      },
      "source": [
        "## Before you begin\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).\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 API, Cloud Resource Manager API](https://console.cloud.google.com/flows/enableapi?apiid=aiplatform.googleapis.com,cloudresourcemanager.googleapis.com).\n",
        "\n",
        "4. If you are running this notebook locally, you need to install the [Cloud SDK](https://cloud.google.com/sdk)."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "project_id"
      },
      "source": [
        "#### Set your project ID\n",
        "\n",
        "**If you don't know your project ID**, try the following:\n",
        "* Run `gcloud config list`.\n",
        "* Run `gcloud projects list`.\n",
        "* See the support page: [Locate the project ID](https://support.google.com/googleapi/answer/7014113)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "wsePm9c4jmpT"
      },
      "outputs": [],
      "source": [
        "PROJECT_ID = \"[your-project-id]\"  # @param {type:\"string\"}\n",
        "\n",
        "# Set the project id\n",
        "! gcloud config set project {PROJECT_ID}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "a54f9d7c1876"
      },
      "source": [
        "#### Region\n",
        "\n",
        "You can also change the `REGION` variable used by Vertex AI. Learn more about [Vertex AI regions](https://cloud.google.com/vertex-ai/docs/general/locations)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "3aaadaaf9b30"
      },
      "outputs": [],
      "source": [
        "REGION = \"us-central1\"  # @param {type: \"string\"}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "5c0404984792"
      },
      "source": [
        "### Authenticate your Google Cloud account\n",
        "\n",
        "Depending on your Jupyter environment, you may have to manually authenticate. Follow the relevant instructions below.\n",
        "\n",
        "**1. Vertex AI Workbench**\n",
        "* Do nothing as you are already authenticated.\n",
        "\n",
        "**2. Local JupyterLab instance, uncomment and run:**"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Nt8cEM2GjmpU"
      },
      "outputs": [],
      "source": [
        "# ! gcloud auth login"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "XUSL_JcpjmpU"
      },
      "source": [
        "**3. Colab, uncomment and run:**"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "_2zemfGvjmpU"
      },
      "outputs": [],
      "source": [
        "# from google.colab import auth\n",
        "# auth.authenticate_user()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "TCPJ38n7jmpU"
      },
      "source": [
        "**4. Service account or other**\n",
        "* See how to grant Cloud Storage permissions to your service account at https://cloud.google.com/storage/docs/gsutil/commands/iam#ch-examples."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "bucket:custom"
      },
      "source": [
        "### Create a Cloud Storage bucket\n",
        "\n",
        "Create a storage bucket to store intermediate artifacts such as datasets.\n",
        "\n",
        "When you submit a training job using the Cloud SDK, you upload a Python package\n",
        "containing your training code to a Cloud Storage bucket. Vertex AI runs\n",
        "the code from this package. In this tutorial, Vertex AI also saves the\n",
        "trained model that results from your job in the same bucket. Using this model artifact, you can then\n",
        "create Vertex AI Model resource and use for prediction."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "bucket"
      },
      "outputs": [],
      "source": [
        "BUCKET_URI = f\"gs://your-bucket-name-{PROJECT_ID}-unique\"  # @param {type:\"string\"}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "create_bucket"
      },
      "source": [
        "**If your bucket doesn't already exist**: Run the following cell to create your Cloud Storage bucket."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Oz8J0vmSlugt"
      },
      "outputs": [],
      "source": [
        "! gsutil mb -l $REGION -p $PROJECT_ID $BUCKET_URI"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "hwXxa4Qgnh4Y"
      },
      "source": [
        "## Setup service account and permissions"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "7qXFUiHLoFRw"
      },
      "source": [
        "A service account is used to create custom training job. If you do not want to use your project's Compute Engine service account, set SERVICE_ACCOUNT to another service account ID. You can create a service account by following the [instruction](https://cloud.google.com/iam/docs/creating-managing-service-accounts#creating)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "mpKjfsXumuNV"
      },
      "outputs": [],
      "source": [
        "SERVICE_ACCOUNT = \"[your-service-account]\"  # @param {type:\"string\"}"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "QbotjwtqKcX4"
      },
      "outputs": [],
      "source": [
        "import sys\n",
        "\n",
        "IS_COLAB = \"google.colab\" in sys.modules\n",
        "if (\n",
        "    SERVICE_ACCOUNT == \"\"\n",
        "    or SERVICE_ACCOUNT is None\n",
        "    or SERVICE_ACCOUNT == \"[your-service-account]\"\n",
        "):\n",
        "    # Get your service account from gcloud\n",
        "    if not IS_COLAB:\n",
        "        shell_output = ! gcloud auth list 2>/dev/null\n",
        "        SERVICE_ACCOUNT = shell_output[2].replace(\"*\", \"\").strip()\n",
        "\n",
        "    else:  # IS_COLAB:\n",
        "        shell_output = ! gcloud projects describe  $PROJECT_ID\n",
        "        project_number = shell_output[-1].split(\":\")[1].strip().replace(\"'\", \"\")\n",
        "        SERVICE_ACCOUNT = f\"{project_number}-compute@developer.gserviceaccount.com\"\n",
        "\n",
        "    print(\"Service Account:\", SERVICE_ACCOUNT)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "UlDhuciOt5vo"
      },
      "outputs": [],
      "source": [
        "# Grant Cloud Storage permission.\n",
        "! gcloud projects add-iam-policy-binding $PROJECT_ID \\\n",
        "   --member=serviceAccount:$SERVICE_ACCOUNT \\\n",
        "   --role=roles/storage.admin \\\n",
        "   --quiet"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "lTKVB71soRyr"
      },
      "outputs": [],
      "source": [
        "# Grant AI Platform permission.\n",
        "! gcloud projects add-iam-policy-binding $PROJECT_ID \\\n",
        "   --member=serviceAccount:$SERVICE_ACCOUNT \\\n",
        "   --role=roles/aiplatform.user \\\n",
        "   --quiet"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "KIPcg_Xhwvsn"
      },
      "source": [
        "### Import libraries"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "p4w8c1pHw2Yt"
      },
      "outputs": [],
      "source": [
        "import google.cloud.aiplatform as aiplatform"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "cVmlv9sRbCSs"
      },
      "source": [
        "### Initialize Vertex AI SDK for Python\n",
        "Initialize the Vertex AI SDK for Python for your project and corresponding bucket."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "JmEsq1fda_1N"
      },
      "outputs": [],
      "source": [
        "aiplatform.init(project=PROJECT_ID, location=REGION, staging_bucket=BUCKET_URI)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "dR2mOCllvlqN"
      },
      "source": [
        "## Write your training code\n",
        "Your training code must be configured to write TensorBoard logs to the Cloud Storage bucket, the location of which the Vertex AI Training service will automatically make available via a predefined environment variable `AIP_TENSORBOARD_LOG_DIR`.\n",
        "\n",
        "This can usually be done by providing `os.environ['AIP_TENSORBOARD_LOG_DIR']` as the log directory to the open source TensorBoard log writing APIs.\n",
        "\n",
        "For example, in TensorFlow 2.x, you can use following code to create a `tensorboard_callback`:\n",
        "```\n",
        "tensorboard_callback = tf.keras.callbacks.TensorBoard(\n",
        "  log_dir=os.environ['AIP_TENSORBOARD_LOG_DIR'],\n",
        "  histogram_freq=1)\n",
        "```\n",
        "\n",
        "`AIP_TENSORBOARD_LOG_DIR` will be in the `BASE_OUTPUT_DIR` that you provided below when creating the custom training job.\n",
        "\n",
        "We'll use the following sample code as an example:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "balswZevBQt_"
      },
      "outputs": [],
      "source": [
        "# Download the sample code\n",
        "! gsutil cp gs://cloud-samples-data/ai-platform/hello-custom/hello-custom-sample-v1.tar.gz - | tar -xzv\n",
        "%cd hello-custom-sample/"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "JGEJkNvHBmc8"
      },
      "outputs": [],
      "source": [
        "# The training code we want to edit is:\n",
        "! cat trainer/task.py"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "vjdLIqYyDZFS"
      },
      "source": [
        "In `trainer/task.py`, create the `tensorboard_callback` and add the callback to `model.fit(...)`\n",
        "\n",
        "Sample code: \n",
        "```\n",
        "# previous things\n",
        "model.compile(...)\n",
        "\n",
        "tensorboard_callback = tf.keras.callbacks.TensorBoard(\n",
        "  log_dir=os.environ['AIP_TENSORBOARD_LOG_DIR'],\n",
        "  histogram_freq=1)\n",
        "  \n",
        "model.fit(dataset, epochs=10, callbacks=[tensorboard_callback])\n",
        "```\n",
        "\n",
        "Update your `trainer/task.py` with `tensorboard_callback`. You can use the following sample code."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "HvI4Y5edurEp"
      },
      "outputs": [],
      "source": [
        "%%writefile trainer/task.py\n",
        "\n",
        "import logging\n",
        "import os\n",
        "\n",
        "import tensorflow as tf\n",
        "import tensorflow_datasets as tfds\n",
        "\n",
        "IMG_WIDTH = 128\n",
        "\n",
        "\n",
        "def normalize_img(image):\n",
        "    \"\"\"Normalizes image.\n",
        "\n",
        "    * Resizes image to IMG_WIDTH x IMG_WIDTH pixels\n",
        "    * Casts values from `uint8` to `float32`\n",
        "    * Scales values from [0, 255] to [0, 1]\n",
        "\n",
        "    Returns:\n",
        "      A tensor with shape (IMG_WIDTH, IMG_WIDTH, 3). (3 color channels)\n",
        "    \"\"\"\n",
        "    image = tf.image.resize_with_pad(image, IMG_WIDTH, IMG_WIDTH)\n",
        "    return image / 255.\n",
        "\n",
        "\n",
        "def normalize_img_and_label(image, label):\n",
        "    \"\"\"Normalizes image and label.\n",
        "\n",
        "    * Performs normalize_img on image\n",
        "    * Passes through label unchanged\n",
        "\n",
        "    Returns:\n",
        "      Tuple (image, label) where\n",
        "      * image is a tensor with shape (IMG_WIDTH, IMG_WIDTH, 3). (3 color\n",
        "        channels)\n",
        "      * label is an unchanged integer [0, 4] representing flower type\n",
        "    \"\"\"\n",
        "    return normalize_img(image), label\n",
        "\n",
        "\n",
        "if 'AIP_MODEL_DIR' not in os.environ:\n",
        "    raise KeyError(\n",
        "        'The `AIP_MODEL_DIR` environment variable has not been' +\n",
        "        'set. See https://cloud.google.com/ai-platform-unified/docs/tutorials/image-recognition-custom/training'\n",
        "    )\n",
        "output_directory = os.environ['AIP_MODEL_DIR']\n",
        "\n",
        "logging.info('Loading and preprocessing data ...')\n",
        "dataset = tfds.load('tf_flowers:3.*.*',\n",
        "                    split='train',\n",
        "                    try_gcs=True,\n",
        "                    shuffle_files=True,\n",
        "                    as_supervised=True)\n",
        "dataset = dataset.map(normalize_img_and_label,\n",
        "                      num_parallel_calls=tf.data.experimental.AUTOTUNE)\n",
        "dataset = dataset.cache()\n",
        "dataset = dataset.shuffle(1000)\n",
        "dataset = dataset.batch(128)\n",
        "dataset = dataset.prefetch(tf.data.experimental.AUTOTUNE)\n",
        "\n",
        "logging.info('Creating and training model ...')\n",
        "model = tf.keras.Sequential([\n",
        "    tf.keras.layers.Conv2D(16,\n",
        "                           3,\n",
        "                           padding='same',\n",
        "                           activation='relu',\n",
        "                           input_shape=(IMG_WIDTH, IMG_WIDTH, 3)),\n",
        "    tf.keras.layers.MaxPooling2D(),\n",
        "    tf.keras.layers.Conv2D(32, 3, padding='same', activation='relu'),\n",
        "    tf.keras.layers.MaxPooling2D(),\n",
        "    tf.keras.layers.Conv2D(64, 3, padding='same', activation='relu'),\n",
        "    tf.keras.layers.MaxPooling2D(),\n",
        "    tf.keras.layers.Flatten(),\n",
        "    tf.keras.layers.Dense(512, activation=\"relu\"),\n",
        "    tf.keras.layers.Dense(5)  # 5 classes\n",
        "])\n",
        "model.compile(\n",
        "    optimizer='adam',\n",
        "    loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),\n",
        "    metrics=['accuracy'])\n",
        "\n",
        "### Create a TensorBoard call back and write to the gcs path provided by AIP_TENSORBOARD_LOG_DIR\n",
        "tensorboard_callback = tf.keras.callbacks.TensorBoard(\n",
        "  log_dir=os.environ['AIP_TENSORBOARD_LOG_DIR'],\n",
        "  histogram_freq=1)\n",
        "\n",
        "### Train the model with tensorboard_callback\n",
        "model.fit(dataset, epochs=14, callbacks=[tensorboard_callback])\n",
        "\n",
        "logging.info(f'Exporting SavedModel to: {output_directory}')\n",
        "# Add softmax layer for intepretability\n",
        "probability_model = tf.keras.Sequential([model, tf.keras.layers.Softmax()])\n",
        "probability_model.save(output_directory)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "L067Jw_QFcZ3"
      },
      "source": [
        "## Upload your training code to Cloud Storage\n",
        "\n",
        "You must package the training code as a source distribution and upload it to Cloud Storage in order for Vertex AI to run the code in a custom training pipeline.\n",
        "\n",
        "1. Run the following commands to create a source distribution in the gzipped tarball format. The command uses the `setup.py` file included in the sample code."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "bNEr950xGbGT"
      },
      "outputs": [],
      "source": [
        "! python3 setup.py sdist --formats=gztar"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "a7Gt9WPFG9V7"
      },
      "source": [
        "2. Run the following command to upload the source distribution you just created, `dist/hello-custom-training-3.0.tar.gz`, to the Cloud Storage bucket that you created before."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "wTv1gReWJqq0"
      },
      "outputs": [],
      "source": [
        "GCS_BUCKET_TRAINING = f\"{BUCKET_URI}/data/\"\n",
        "! gsutil cp dist/hello-custom-training-3.0.tar.gz {GCS_BUCKET_TRAINING}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "IaQjIPvuKLwW"
      },
      "source": [
        "## Create a custom training job"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "svUGBOow_Obj"
      },
      "source": [
        "Create a TensorBoard instance to be used by the custom training job."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "OAe1xJeS_X3F"
      },
      "outputs": [],
      "source": [
        "TENSORBOARD_NAME = \"[your-tensorboard-name]\"  # @param {type:\"string\"}\n",
        "\n",
        "if (\n",
        "    TENSORBOARD_NAME == \"\"\n",
        "    or TENSORBOARD_NAME is None\n",
        "    or TENSORBOARD_NAME == \"[your-tensorboard-name]\"\n",
        "):\n",
        "    TENSORBOARD_NAME = PROJECT_ID + \"-tb\"\n",
        "\n",
        "tensorboard = aiplatform.Tensorboard.create(\n",
        "    display_name=TENSORBOARD_NAME, project=PROJECT_ID, location=REGION\n",
        ")\n",
        "TENSORBOARD_RESOURCE_NAME = tensorboard.gca_resource.name\n",
        "print(\"TensorBoard resource name:\", TENSORBOARD_RESOURCE_NAME)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "mudxBDal_a_k"
      },
      "source": [
        "Run following example request to create your own custom training job and stream the training results to TensorBoard."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "le-kJx895rqN"
      },
      "outputs": [],
      "source": [
        "JOB_NAME = \"tensorboard-example-job\"\n",
        "GCS_BUCKET_OUTPUT = BUCKET_URI\n",
        "BASE_OUTPUT_DIR = \"{}/{}\".format(GCS_BUCKET_OUTPUT, JOB_NAME)\n",
        "\n",
        "job = aiplatform.CustomPythonPackageTrainingJob(\n",
        "    display_name=JOB_NAME,\n",
        "    python_module_name=\"trainer.task\",\n",
        "    container_uri=\"us-docker.pkg.dev/vertex-ai/training/tf-cpu.2-8:latest\",\n",
        "    python_package_gcs_uri=f\"{GCS_BUCKET_TRAINING}hello-custom-training-3.0.tar.gz\",\n",
        "    project=PROJECT_ID,\n",
        "    location=REGION,\n",
        "    staging_bucket=BASE_OUTPUT_DIR,\n",
        ")\n",
        "\n",
        "job.run(\n",
        "    service_account=SERVICE_ACCOUNT,\n",
        "    tensorboard=TENSORBOARD_RESOURCE_NAME,\n",
        "    machine_type=\"n1-standard-8\",\n",
        "    replica_count=1,\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "VfMsn_RnEtnj"
      },
      "source": [
        "In Google Cloud Console, you can monitor your training job at Vertex AI > Training > Custom Jobs. In each custom training job, near real time updated TensorBoard is available at `OPEN TENSORBOARD` button. Learn more see [View Vertex AI TensorBoard data](https://cloud.google.com/vertex-ai/docs/experiments/tensorboard-view)."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "TFEriiywMZga"
      },
      "source": [
        "## Cleaning up\n",
        "\n",
        "To clean up all Google Cloud resources used in this project, you can [delete the Google Cloud project](https://cloud.google.com/resource-manager/docs/creating-managing-projects#shutting_down_projects) you used for the tutorial.\n",
        "\n",
        "Otherwise, **if you created the individual resources in the notebook** you can delete them as follow:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "O1Je2JZTMoMg"
      },
      "outputs": [],
      "source": [
        "# Delete GCS bucket.\n",
        "! gsutil -m rm -r {BUCKET_URI}\n",
        "\n",
        "# Delete TensorBoard instance.\n",
        "! gcloud ai tensorboards delete {TENSORBOARD_RESOURCE_NAME}\n",
        "\n",
        "# Delete custom job.\n",
        "job.delete()"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "collapsed_sections": [],
      "name": "tensorboard_custom_training_with_prebuilt_container.ipynb",
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}