notebooks/official/migration/sdk-hyperparameter-tuning.ipynb

{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "copyright"
      },
      "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": "title:migration,new"
      },
      "source": [
        "# Vertex AI Migration: Hyperparameter Tuning\n",
        "\n",
        "<table align=\"left\">\n",
        "\n",
        "  <td>\n",
        "    <a href=\"https://colab.research.google.com/github/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/official/migration/sdk-hyperparameter-tuning.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/migration/sdk-hyperparameter-tuning.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/migration/sdk-hyperparameter-tuning.ipynb\" target='_blank'>\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": "7a8a13b86a8b"
      },
      "source": [
        "## Overview\n",
        "\n",
        "\n",
        "This tutorial demonstrates how to use the Vertex AI SDK for Python to tune hyperparameters in a custom tabular classification TensorFlow model.\n",
        "\n",
        "Learn more about [Migrate to Vertex AI](https://cloud.google.com/vertex-ai/docs/training/hyperparameter-tuning-overview) and [Custom training](https://cloud.google.com/vertex-ai/docs/training/custom-training)."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "618cfedf829a"
      },
      "source": [
        "### Objective\n",
        "\n",
        "In this tutorial, you learn to use `Vertex AI Hyperparameter` to create and tune a custom trained model.\n",
        "\n",
        "You learn how to create and tune a custom-trained model from a Python script in a Docker container using the Vertex AI SDK for Python.\n",
        "\n",
        "This tutorial uses the following Google Cloud ML services:\n",
        "\n",
        "- `Vertex AI Training`\n",
        "- `Vertex AI Hyperparameter Tuning`\n",
        "\n",
        "The steps performed include:\n",
        "\n",
        "- Create a `Vertex AI` hyperparameter tuning job for training a TensorFlow model."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "dataset:custom,boston,lrg"
      },
      "source": [
        "### Dataset\n",
        "\n",
        "The dataset used for this tutorial is the [Boston Housing Prices dataset](https://www.cs.toronto.edu/~delve/data/boston/bostonDetail.html). The version of the dataset you use in 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": "install_aip:mbsdk"
      },
      "source": [
        "## Installation\n",
        "\n",
        "Install the packages required for executing this notebook."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "install_aip:mbsdk"
      },
      "outputs": [],
      "source": [
        "! pip3 install --upgrade --quiet google-cloud-aiplatform \\\n",
        "                                 google-cloud-storage"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "restart"
      },
      "source": [
        "### Colab only: Uncomment the following cell to restart the kernel"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "D-ZBOjErv5mM"
      },
      "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": "before_you_begin:nogpu"
      },
      "source": [
        "## Before you begin"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "before_you_begin:nogpu"
      },
      "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": "set_project_id"
      },
      "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": "region"
      },
      "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": "region"
      },
      "outputs": [],
      "source": [
        "REGION = \"us-central1\"  # @param {type: \"string\"}"
      ]
    },
    {
      "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",
        "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": "ce6043da7b33"
      },
      "outputs": [],
      "source": [
        "# ! gcloud auth login"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "0367eac06a10"
      },
      "source": [
        "**3. Colab, uncomment and run:**"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "21ad4dbb4a61"
      },
      "outputs": [],
      "source": [
        "# from google.colab import auth\n",
        "# auth.authenticate_user()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "c13224697bfb"
      },
      "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:mbsdk"
      },
      "source": [
        "### Create a Cloud Storage bucket\n",
        "\n",
        "Create a storage bucket to store intermediate artifacts such as datasets."
      ]
    },
    {
      "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": [
        "**Only 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": "create_bucket"
      },
      "outputs": [],
      "source": [
        "! gsutil mb -l {REGION} {BUCKET_URI}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "setup_vars"
      },
      "source": [
        "### Set up variables\n",
        "\n",
        "Next, set up some variables used throughout the tutorial.\n",
        "### Import libraries and define constants"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "import_aip:mbsdk"
      },
      "outputs": [],
      "source": [
        "import google.cloud.aiplatform as aip"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "init_aip:mbsdk"
      },
      "source": [
        "## Initialize Vertex SDK for Python\n",
        "\n",
        "Initialize the Vertex SDK for Python for your project and corresponding bucket."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "init_aip:mbsdk"
      },
      "outputs": [],
      "source": [
        "aip.init(project=PROJECT_ID, staging_bucket=BUCKET_URI)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "accelerators:training,cpu,prediction,cpu,mbsdk"
      },
      "source": [
        "#### Set hardware accelerators\n",
        "\n",
        "You can set hardware accelerators for training and prediction.\n",
        "\n",
        "Set the variables `TRAIN_GPU/TRAIN_NGPU` and `DEPLOY_GPU/DEPLOY_NGPU` to use a container image supporting a GPU and the number of GPUs allocated to the virtual machine (VM) instance. For example, to use a GPU container image with 4 Nvidia Telsa K80 GPUs allocated to each VM, you would specify:\n",
        "\n",
        "    (aip.gapic.AcceleratorType.NVIDIA_TESLA_K80, 4)\n",
        "\n",
        "\n",
        "Otherwise specify `(None, None)` to use a container image to run on a CPU.\n",
        "\n",
        "Learn more [here](https://cloud.google.com/vertex-ai/docs/general/locations#accelerators) hardware accelerator support for your region"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "accelerators:training,cpu,prediction,cpu,mbsdk"
      },
      "outputs": [],
      "source": [
        "TRAIN_GPU, TRAIN_NGPU = (None, None)\n",
        "\n",
        "DEPLOY_GPU, DEPLOY_NGPU = (None, None)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "container:training,prediction"
      },
      "source": [
        "#### Set pre-built containers\n",
        "\n",
        "Set the pre-built Docker container image for training and prediction.\n",
        "\n",
        "\n",
        "For the latest list, see [Pre-built containers for training](https://cloud.google.com/ai-platform-unified/docs/training/pre-built-containers).\n",
        "\n",
        "\n",
        "For the latest list, see [Pre-built containers for prediction](https://cloud.google.com/ai-platform-unified/docs/predictions/pre-built-containers)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "container:training,prediction"
      },
      "outputs": [],
      "source": [
        "TF = \"2-5\"\n",
        "\n",
        "if TRAIN_GPU:\n",
        "    TRAIN_VERSION = \"tf-gpu.{}\".format(TF)\n",
        "else:\n",
        "    TRAIN_VERSION = \"tf-cpu.{}\".format(TF)\n",
        "if DEPLOY_GPU:\n",
        "    DEPLOY_VERSION = \"tf2-gpu.{}\".format(TF)\n",
        "else:\n",
        "    DEPLOY_VERSION = \"tf2-cpu.{}\".format(TF)\n",
        "\n",
        "TRAIN_IMAGE = \"us-docker.pkg.dev/vertex-ai/training/{}:latest\".format(TRAIN_VERSION)\n",
        "DEPLOY_IMAGE = \"us-docker.pkg.dev/vertex-ai/prediction/{}:latest\".format(DEPLOY_VERSION)\n",
        "\n",
        "print(\"Training:\", TRAIN_IMAGE, TRAIN_GPU, TRAIN_NGPU)\n",
        "print(\"Deployment:\", DEPLOY_IMAGE, DEPLOY_GPU, DEPLOY_NGPU)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "machine:training,prediction"
      },
      "source": [
        "#### Set machine type\n",
        "\n",
        "Next, set the machine type to use for training and prediction.\n",
        "\n",
        "- Set the variables `TRAIN_COMPUTE` and `DEPLOY_COMPUTE` to configure  the compute resources for the VMs you use for for training and prediction.\n",
        " - `machine type`\n",
        "     - `n1-standard`: 3.75GB of memory per vCPU.\n",
        "     - `n1-highmem`: 6.5GB of memory per vCPU\n",
        "     - `n1-highcpu`: 0.9 GB of memory per vCPU\n",
        " - `vCPUs`: number of \\[2, 4, 8, 16, 32, 64, 96 \\]\n",
        "\n",
        "*Note: The following is not supported for training:*\n",
        "\n",
        " - `standard`: 2 vCPUs\n",
        " - `highcpu`: 2, 4 and 8 vCPUs\n",
        "\n",
        "*Note: You may also use n2 and e2 machine types for training and deployment, but they do not support GPUs*."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "machine:training,prediction"
      },
      "outputs": [],
      "source": [
        "MACHINE_TYPE = \"n1-standard-4\"\n",
        "\n",
        "TRAIN_COMPUTE = MACHINE_TYPE\n",
        "print(\"Train machine type\", TRAIN_COMPUTE)\n",
        "\n",
        "DEPLOY_COMPUTE = MACHINE_TYPE\n",
        "print(\"Deploy machine type\", DEPLOY_COMPUTE)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "examine_training_package"
      },
      "source": [
        "### Examine the training package\n",
        "\n",
        "#### Package layout\n",
        "\n",
        "Before you start the training, you look at how a Python package is assembled for a custom training job. When unarchived, the package contains the following directory/file layout.\n",
        "\n",
        "- PKG-INFO\n",
        "- README.md\n",
        "- setup.cfg\n",
        "- setup.py\n",
        "- trainer\n",
        "  - \\_\\_init\\_\\_.py\n",
        "  - task.py\n",
        "\n",
        "The files `setup.cfg` and `setup.py` are the instructions for installing the package into the operating environment of the Docker image.\n",
        "\n",
        "The file `trainer/task.py` is the Python script for executing the custom training job. *Note*, when we referred to it in the worker pool specification, we replace the directory slash with a dot (`trainer.task`) and dropped the file suffix (`.py`).\n",
        "\n",
        "#### Package Assembly\n",
        "\n",
        "In the following cells, you assemble the training package."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "examine_training_package"
      },
      "outputs": [],
      "source": [
        "# Make folder for Python training script\n",
        "! rm -rf custom\n",
        "! mkdir custom\n",
        "\n",
        "# Add package information\n",
        "! touch custom/README.md\n",
        "\n",
        "setup_cfg = \"[egg_info]\\n\\ntag_build =\\n\\ntag_date = 0\"\n",
        "! echo \"$setup_cfg\" > custom/setup.cfg\n",
        "\n",
        "setup_py = \"import setuptools\\n\\nsetuptools.setup(\\n\\n    install_requires=[\\n\\n        'tensorflow_datasets==1.3.0',\\n\\n    ],\\n\\n    packages=setuptools.find_packages())\"\n",
        "! echo \"$setup_py\" > custom/setup.py\n",
        "\n",
        "pkg_info = \"Metadata-Version: 1.0\\n\\nName: Boston Housing tabular regression\\n\\nVersion: 0.0.0\\n\\nSummary: Demostration training script\\n\\nHome-page: www.google.com\\n\\nAuthor: Google\\n\\nAuthor-email: aferlitsch@google.com\\n\\nLicense: Public\\n\\nDescription: Demo\\n\\nPlatform: Vertex\"\n",
        "! echo \"$pkg_info\" > custom/PKG-INFO\n",
        "\n",
        "# Make the training subfolder\n",
        "! mkdir custom/trainer\n",
        "! touch custom/trainer/__init__.py"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "taskpy_contents:hpt,boston"
      },
      "source": [
        "#### Task.py contents\n",
        "\n",
        "In the next cell, you write the contents of the hyperparameter tuning script task.py. I won't go into detail, it's just there for you to browse. In summary:\n",
        "\n",
        "- Parse the command line arguments for the hyperparameter settings for the current trial.\n",
        " - Get the directory where to save the model artifacts from the command line (`--model_dir`), and if not specified, then from the environment variable `AIP_MODEL_DIR`.\n",
        "- Download and preprocess the Boston Housing dataset.\n",
        "- Build a DNN model.\n",
        "- The number of units per dense layer and learning rate hyperparameter values are used during the build and compile of the model.\n",
        "- A definition of a callback `HPTCallback` which obtains the validation loss at the end of each epoch (`on_epoch_end()`) and reports it to the hyperparameter tuning service using `hpt.report_hyperparameter_tuning_metric()`.\n",
        "- Train the model with the `fit()` method and specify a callback which report the validation loss back to the hyperparameter tuning service."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "taskpy_contents:hpt,boston"
      },
      "outputs": [],
      "source": [
        "%%writefile custom/trainer/task.py\n",
        "# Custom Training for Boston Housing\n",
        "\n",
        "import tensorflow_datasets as tfds\n",
        "import tensorflow as tf\n",
        "from tensorflow.python.client import device_lib\n",
        "from hypertune import HyperTune\n",
        "import numpy as np\n",
        "import argparse\n",
        "import os\n",
        "import sys\n",
        "tfds.disable_progress_bar()\n",
        "\n",
        "parser = argparse.ArgumentParser()\n",
        "parser.add_argument('--model-dir', dest='model_dir',\n",
        "                    default=os.getenv('AIP_MODEL_DIR'), type=str, help='Model dir.')\n",
        "parser.add_argument('--lr', dest='lr',\n",
        "                    default=0.001, type=float,\n",
        "                    help='Learning rate.')\n",
        "parser.add_argument('--decay', dest='decay',\n",
        "                    default=0.98, type=float,\n",
        "                    help='Decay rate')\n",
        "parser.add_argument('--units', dest='units',\n",
        "                    default=64, type=int,\n",
        "                    help='Number of units.')\n",
        "parser.add_argument('--epochs', dest='epochs',\n",
        "                    default=20, type=int,\n",
        "                    help='Number of epochs.')\n",
        "parser.add_argument('--steps', dest='steps',\n",
        "                    default=200, type=int,\n",
        "                    help='Number of steps per epoch.')\n",
        "parser.add_argument('--param-file', dest='param_file',\n",
        "                    default='/tmp/param.txt', type=str,\n",
        "                    help='Output file for parameters')\n",
        "parser.add_argument('--distribute', dest='distribute', type=str, default='single',\n",
        "                    help='distributed training strategy')\n",
        "args = parser.parse_args()\n",
        "\n",
        "print('Python Version = {}'.format(sys.version))\n",
        "print('TensorFlow Version = {}'.format(tf.__version__))\n",
        "print('TF_CONFIG = {}'.format(os.environ.get('TF_CONFIG', 'Not found')))\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",
        "  for _ in range(13):\n",
        "    x_train[_], max = scale(x_train[_])\n",
        "    x_test[_], _ = scale(x_test[_])\n",
        "    params.append(max)\n",
        "\n",
        "  # store the normalization (max) value for each feature\n",
        "  with tf.io.gfile.GFile(args.param_file, 'w') as f:\n",
        "    f.write(str(params))\n",
        "  return (x_train, y_train), (x_test, y_test)\n",
        "\n",
        "# Build the Keras model\n",
        "def build_and_compile_dnn_model():\n",
        "  model = tf.keras.Sequential([\n",
        "      tf.keras.layers.Dense(args.units, activation='relu', input_shape=(13,)),\n",
        "      tf.keras.layers.Dense(args.units, activation='relu'),\n",
        "      tf.keras.layers.Dense(1, activation='linear')\n",
        "  ])\n",
        "  model.compile(\n",
        "      loss='mse',\n",
        "      optimizer=tf.keras.optimizers.RMSprop(learning_rate=args.lr, decay=args.decay))\n",
        "  return model\n",
        "\n",
        "\n",
        "model = build_and_compile_dnn_model()\n",
        "\n",
        "# Instantiate the HyperTune reporting object\n",
        "hpt = HyperTune()\n",
        "\n",
        "# Reporting callback\n",
        "class HPTCallback(tf.keras.callbacks.Callback):\n",
        "\n",
        "    def on_epoch_end(self, epoch, logs=None):\n",
        "        global hpt\n",
        "        hpt.report_hyperparameter_tuning_metric(\n",
        "        hyperparameter_metric_tag='val_loss',\n",
        "        metric_value=logs['val_loss'],\n",
        "        global_step=epoch)\n",
        "\n",
        "# Train the model\n",
        "BATCH_SIZE = 16\n",
        "(x_train, y_train), (x_test, y_test) = make_dataset()\n",
        "model.fit(x_train, y_train, epochs=args.epochs, batch_size=BATCH_SIZE, validation_split=0.1, callbacks=[HPTCallback()])\n",
        "model.save(args.model_dir)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "tarball_training_script"
      },
      "source": [
        "#### Store training script on your Cloud Storage bucket\n",
        "\n",
        "Next, you package the training folder into a compressed tar ball, and then store it in your Cloud Storage bucket."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "tarball_training_script"
      },
      "outputs": [],
      "source": [
        "! rm -f custom.tar custom.tar.gz\n",
        "! tar cvf custom.tar custom\n",
        "! gzip custom.tar\n",
        "! gsutil cp custom.tar.gz $BUCKET_URI/trainer_boston.tar.gz"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "train_a_model:migration"
      },
      "source": [
        "## Train a model"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "hyperparameter_tuning:migration,new,mbsdk"
      },
      "source": [
        "### [training.using-hyperparamter-tuning](https://cloud.google.com/vertex-ai/docs/training/using-hyperparameter-tuning)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "train_custom_job_machine_specification"
      },
      "source": [
        "### Prepare your machine specification\n",
        "\n",
        "Now define the machine specification for your custom training job. This tells Vertex what type of machine instance to provision for the training.\n",
        "  - `machine_type`: The type of GCP instance to provision -- e.g., n1-standard-8.\n",
        "  - `accelerator_type`: The type, if any, of hardware accelerator. In this tutorial if you previously set the variable `TRAIN_GPU != None`, you are using a GPU; otherwise you use a CPU.\n",
        "  - `accelerator_count`: The number of accelerators."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "train_custom_job_machine_specification"
      },
      "outputs": [],
      "source": [
        "if TRAIN_GPU:\n",
        "    machine_spec = {\n",
        "        \"machine_type\": TRAIN_COMPUTE,\n",
        "        \"accelerator_type\": TRAIN_GPU,\n",
        "        \"accelerator_count\": TRAIN_NGPU,\n",
        "    }\n",
        "else:\n",
        "    machine_spec = {\"machine_type\": TRAIN_COMPUTE, \"accelerator_count\": 0}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "train_custom_job_disk_specification"
      },
      "source": [
        "### Prepare your disk specification\n",
        "\n",
        "(optional) Now define the disk specification for your custom training job. This tells Vertex what type and size of disk to provision in each machine instance for the training.\n",
        "\n",
        "  - `boot_disk_type`: Either SSD or Standard. SSD is faster, and Standard is less expensive. Defaults to SSD.\n",
        "  - `boot_disk_size_gb`: Size of disk in GB."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "train_custom_job_disk_specification"
      },
      "outputs": [],
      "source": [
        "DISK_TYPE = \"pd-ssd\"  # [ pd-ssd, pd-standard]\n",
        "DISK_SIZE = 200  # GB\n",
        "\n",
        "disk_spec = {\"boot_disk_type\": DISK_TYPE, \"boot_disk_size_gb\": DISK_SIZE}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "train_custom_job_worker_pool_specification:prebuilt_container"
      },
      "source": [
        "### Define the worker pool specification\n",
        "\n",
        "Next, you define the worker pool specification for your custom training job. The worker pool specification  consist of the following:\n",
        "\n",
        "- `replica_count`: The number of instances to provision of this machine type.\n",
        "- `machine_spec`: The hardware specification.\n",
        "- `disk_spec` : (optional) The disk storage specification.\n",
        "\n",
        "- `python_package`: The Python training package to install on the VM instance(s) and which Python module to invoke, along with command line arguments for the Python module.\n",
        "\n",
        "Let's dive deeper now into the python package specification:\n",
        "\n",
        "-`executor_image_spec`: This is the docker image which is configured for your custom training job.\n",
        "\n",
        "-`package_uris`: This is a list of the locations (URIs) of your python training packages to install on the provisioned instance. The locations need to be in a Cloud Storage bucket. These can be either individual python files or a zip (archive) of an entire package. In the later case, the job service  unzip (unarchive) the contents into the docker image.\n",
        "\n",
        "-`python_module`: The Python module (script) to invoke for running the custom training job. In this example, you  be invoking `trainer.task.py` -- note that it was not neccessary to append the `.py` suffix.\n",
        "\n",
        "-`args`: The command line arguments to pass to the corresponding Pythom module. In this example, you  be setting:\n",
        "  - `\"--model-dir=\" + MODEL_DIR` : The Cloud Storage location where to store the model artifacts. There are two ways to tell the training script where to save the model artifacts:\n",
        "      - direct: You pass the Cloud Storage location as a command line argument to your training script (set variable `DIRECT = True`), or\n",
        "      - indirect: The service passes the Cloud Storage location as the environment variable `AIP_MODEL_DIR` to your training script (set variable `DIRECT = False`). In this case, you tell the service the model artifact location in the job specification.\n",
        "  - `\"--epochs=\" + EPOCHS`: The number of epochs for training.\n",
        "  - `\"--steps=\" + STEPS`: The number of steps (batches) per epoch.\n",
        "  - `\"--distribute=\" + TRAIN_STRATEGY\"` : The training distribution strategy to use for single or distributed training.\n",
        "     - `\"single\"`: single device.\n",
        "     - `\"mirror\"`: all GPU devices on a single compute instance.\n",
        "     - `\"multi\"`: all GPU devices on all compute instances."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "train_custom_job_worker_pool_specification:prebuilt_container"
      },
      "outputs": [],
      "source": [
        "JOB_NAME = \"custom_job_\" + UUID\n",
        "MODEL_DIR = \"{}/{}\".format(BUCKET_URI, JOB_NAME)\n",
        "\n",
        "if not TRAIN_NGPU or TRAIN_NGPU < 2:\n",
        "    TRAIN_STRATEGY = \"single\"\n",
        "else:\n",
        "    TRAIN_STRATEGY = \"mirror\"\n",
        "\n",
        "EPOCHS = 20\n",
        "STEPS = 100\n",
        "\n",
        "DIRECT = True\n",
        "if DIRECT:\n",
        "    CMDARGS = [\n",
        "        \"--model-dir=\" + MODEL_DIR,\n",
        "        \"--epochs=\" + str(EPOCHS),\n",
        "        \"--steps=\" + str(STEPS),\n",
        "        \"--distribute=\" + TRAIN_STRATEGY,\n",
        "    ]\n",
        "else:\n",
        "    CMDARGS = [\n",
        "        \"--epochs=\" + str(EPOCHS),\n",
        "        \"--steps=\" + str(STEPS),\n",
        "        \"--distribute=\" + TRAIN_STRATEGY,\n",
        "    ]\n",
        "\n",
        "worker_pool_spec = [\n",
        "    {\n",
        "        \"replica_count\": 1,\n",
        "        \"machine_spec\": machine_spec,\n",
        "        \"disk_spec\": disk_spec,\n",
        "        \"python_package_spec\": {\n",
        "            \"executor_image_uri\": TRAIN_IMAGE,\n",
        "            \"package_uris\": [BUCKET_URI + \"/trainer_boston.tar.gz\"],\n",
        "            \"python_module\": \"trainer.task\",\n",
        "            \"args\": CMDARGS,\n",
        "        },\n",
        "    }\n",
        "]"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "create_custom_job:migration,new,mbsdk"
      },
      "source": [
        "### [training.create-custom-job](https://cloud.google.com/vertex-ai/docs/training/create-custom-job)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "create_custom_job:mbsdk"
      },
      "source": [
        "## Create a custom job\n",
        "\n",
        "Use the class `CustomJob` to create a custom job, such as for hyperparameter tuning, with the following parameters:\n",
        "\n",
        "- `display_name`: A human readable name for the custom job.\n",
        "- `worker_pool_specs`: The specification for the corresponding VM instances."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "create_custom_job:mbsdk"
      },
      "outputs": [],
      "source": [
        "job = aip.CustomJob(display_name=\"boston_\" + UUID, worker_pool_specs=worker_pool_spec)\n",
        "\n",
        "# print(job)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "create_hpt_job:mbsdk"
      },
      "source": [
        "## Create a hyperparameter tuning job\n",
        "\n",
        "Use the class `HyperparameterTuningJob` to create a hyperparameter tuning job, with the following parameters:\n",
        "\n",
        "- `display_name`: A human readable name for the custom job.\n",
        "- `custom_job`: The worker pool spec from this custom job applies to the CustomJobs created in all the trials.\n",
        "- `metrics_spec`: The metrics to optimize. The dictionary key is the metric_id, which is reported by your training job, and the dictionary value is the optimization goal of the metric('minimize' or 'maximize').\n",
        "- `parameter_spec`: The parameters to optimize. The dictionary key is the metric_id, which is passed into your training job as a command line key word argument, and the dictionary value is the parameter specification of the metric."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "create_hpt_job:mbsdk"
      },
      "outputs": [],
      "source": [
        "from google.cloud.aiplatform import hyperparameter_tuning as hpt\n",
        "\n",
        "hpt_job = aip.HyperparameterTuningJob(\n",
        "    display_name=\"boston_\" + UUID,\n",
        "    custom_job=job,\n",
        "    metric_spec={\n",
        "        \"val_loss\": \"minimize\",\n",
        "    },\n",
        "    parameter_spec={\n",
        "        \"lr\": hpt.DoubleParameterSpec(min=0.001, max=0.1, scale=\"log\"),\n",
        "        \"units\": hpt.IntegerParameterSpec(min=4, max=128, scale=\"linear\"),\n",
        "    },\n",
        "    max_trial_count=6,\n",
        "    parallel_trial_count=1,\n",
        ")\n",
        "\n",
        "# print(hpt_job)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "run_hpt_job:mbsdk"
      },
      "source": [
        "## Run the hyperparameter tuning job\n",
        "\n",
        "Use the `run()` method to execute the hyperparameter tuning job."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "run_hpt_job:mbsdk"
      },
      "outputs": [],
      "source": [
        "hpt_job.run()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "run_hpt_job:mbsdk"
      },
      "source": [
        "*Example output:*\n",
        "\n",
        "    INFO:google.cloud.aiplatform.jobs:Creating HyperparameterTuningJob\n",
        "    INFO:google.cloud.aiplatform.jobs:HyperparameterTuningJob created. Resource name: projects/759209241365/locations/us-central1/hyperparameterTuningJobs/760969798560514048\n",
        "    INFO:google.cloud.aiplatform.jobs:To use this HyperparameterTuningJob in another session:\n",
        "    INFO:google.cloud.aiplatform.jobs:hpt_job = aiplatform.HyperparameterTuningJob.get('projects/759209241365/locations/us-central1/hyperparameterTuningJobs/760969798560514048')\n",
        "    INFO:google.cloud.aiplatform.jobs:View HyperparameterTuningJob:\n",
        "    https://console.cloud.google.com/ai/platform/locations/us-central1/training/760969798560514048?project=759209241365\n",
        "    INFO:google.cloud.aiplatform.jobs:HyperparameterTuningJob projects/759209241365/locations/us-central1/hyperparameterTuningJobs/760969798560514048 current state:\n",
        "    JobState.JOB_STATE_RUNNING\n",
        "    INFO:google.cloud.aiplatform.jobs:HyperparameterTuningJob projects/759209241365/locations/us-central1/hyperparameterTuningJobs/760969798560514048 current state:\n",
        "    JobState.JOB_STATE_RUNNING\n",
        "    ...\n",
        "    INFO:google.cloud.aiplatform.jobs:HyperparameterTuningJob projects/759209241365/locations/us-central1/hyperparameterTuningJobs/760969798560514048 current state:\n",
        "    JobState.JOB_STATE_SUCCEEDED\n",
        "    INFO:google.cloud.aiplatform.jobs:HyperparameterTuningJob run completed. Resource name: projects/759209241365/locations/us-central1/hyperparameterTuningJobs/760969798560514048"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "hpt_trial_results:mbsdk"
      },
      "source": [
        "### Display the hyperparameter tuning job trial results\n",
        "\n",
        "After the hyperparameter tuning job has completed, the property `trials`  return the results for each trial."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "hpt_trial_results:mbsdk"
      },
      "outputs": [],
      "source": [
        "print(hpt_job.trials)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "hpt_trial_results:mbsdk"
      },
      "source": [
        "*Example output:*\n",
        "\n",
        "    [id: \"1\"\n",
        "    state: SUCCEEDED\n",
        "    parameters {\n",
        "      parameter_id: \"lr\"\n",
        "      value {\n",
        "        number_value: 0.010000000000000002\n",
        "      }\n",
        "    }\n",
        "    parameters {\n",
        "      parameter_id: \"units\"\n",
        "      value {\n",
        "        number_value: 66.0\n",
        "      }\n",
        "    }\n",
        "    final_measurement {\n",
        "      step_count: 19\n",
        "      metrics {\n",
        "        metric_id: \"val_loss\"\n",
        "        value: 24.61802691948123\n",
        "      }\n",
        "    }\n",
        "    start_time {\n",
        "      seconds: 1629414344\n",
        "      nanos: 232151761\n",
        "    }\n",
        "    end_time {\n",
        "      seconds: 1629414825\n",
        "    }\n",
        "    , id: \"2\"\n",
        "    state: SUCCEEDED\n",
        "    parameters {\n",
        "      parameter_id: \"lr\"\n",
        "      value {\n",
        "        number_value: 0.0036327633937958217\n",
        "      }\n",
        "    }\n",
        "    parameters {\n",
        "      parameter_id: \"units\"\n",
        "      value {\n",
        "        number_value: 94.0\n",
        "      }\n",
        "    }\n",
        "    final_measurement {\n",
        "      step_count: 19\n",
        "      metrics {\n",
        "        metric_id: \"val_loss\"\n",
        "        value: 25.499705989186356\n",
        "      }\n",
        "    }\n",
        "    start_time {\n",
        "      seconds: 1629414960\n",
        "      nanos: 618333580\n",
        "    }\n",
        "    end_time {\n",
        "      seconds: 1629415021\n",
        "    }\n",
        "    ...\n",
        "    , id: \"6\"\n",
        "    state: SUCCEEDED\n",
        "    parameters {\n",
        "      parameter_id: \"lr\"\n",
        "      value {\n",
        "        number_value: 0.002775175422799751\n",
        "      }\n",
        "    }\n",
        "    parameters {\n",
        "      parameter_id: \"units\"\n",
        "      value {\n",
        "        number_value: 104.0\n",
        "      }\n",
        "    }\n",
        "    final_measurement {\n",
        "      step_count: 7\n",
        "      metrics {\n",
        "        metric_id: \"val_loss\"\n",
        "        value: 24.655450123112377\n",
        "      }\n",
        "    }\n",
        "    start_time {\n",
        "      seconds: 1629415735\n",
        "      nanos: 998711852\n",
        "    }\n",
        "    end_time {\n",
        "      seconds: 1629415766\n",
        "    }\n",
        "    ]"
      ]
    },
    {
      "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": "cleanup:mbsdk"
      },
      "outputs": [],
      "source": [
        "import os\n",
        "\n",
        "delete_bucket = False\n",
        "\n",
        "# Delete the training job\n",
        "try:\n",
        "    job.delete()\n",
        "except Exception as e:\n",
        "    print(e)\n",
        "\n",
        "# Delete the HPT job using the Vertex batch prediction object\n",
        "hpt_job.delete()\n",
        "\n",
        "if delete_bucket or os.getenv(\"IS_TESTING\"):\n",
        "    ! gsutil rm -r $BUCKET_URI\n"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "name": "sdk-hyperparameter-tuning.ipynb",
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}