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",
        "  <td style=\"text-align: center\">\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=\"Google Colaboratory logo\"><br> Open in Colab\n",
        "    </a>\n",
        "  </td>\n",
        "  <td style=\"text-align: center\">\n",
        "    <a href=\"https://console.cloud.google.com/vertex-ai/colab/import/https:%2F%2Fraw.githubusercontent.com%2FGoogleCloudPlatform%2Fvertex-ai-samples%2Fmain%2Fnotebooks%2Fofficial%2Fmigration%2Fsdk-hyperparameter-tuning.ipynb\">\n",
        "      <img src=\"https://cloud.google.com/ml-engine/images/colab-enterprise-logo-32px.png\" alt=\"Google Cloud Colab Enterprise logo\"><br> Open in Colab Enterprise\n",
        "    </a>\n",
        "  </td>\n",
        "  <td style=\"text-align: center\">\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\"><br> View on GitHub\n",
        "    </a>\n",
        "  </td>\n",
        "  <td style=\"text-align: center\">\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\"><br> Open in Vertex AI Workbench\n",
        "     </a>\n",
        "   </td>\n",
        "</table>\n",
        "<br/><br/><br/>\n"
      ]
    },
    {
      "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 [Vertex AI hyperparameter tuning](https://cloud.google.com/vertex-ai/docs/training/hyperparameter-tuning-overview) and [Custom training](https://cloud.google.com/vertex-ai/docs/training/overview)."
      ]
    },
    {
      "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 pricing](https://cloud.google.com/vertex-ai/pricing) and\n",
        "[Cloud Storage pricing](https://cloud.google.com/storage/pricing), and use the \n",
        "[Pricing Calculator](https://cloud.google.com/products/calculator/) to generate a cost estimate based on your projected usage."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "install_aip:mbsdk"
      },
      "source": [
        "## Get started\n",
        "Install Vertex AI SDK for Python and other required packages"
      ]
    },
    {
      "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": [
        "### Restart runtime (Colab only)\n",
        "To use the newly installed packages, you must restart the runtime on Google Colab."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "D-ZBOjErv5mM"
      },
      "outputs": [],
      "source": [
        "import sys\n",
        "\n",
        "if \"google.colab\" in sys.modules:\n",
        "\n",
        "    import IPython\n",
        "\n",
        "    app = IPython.Application.instance()\n",
        "    app.kernel.do_shutdown(True)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "4de1bd77992b"
      },
      "source": [
        "<div class=\"alert alert-block alert-warning\">,\n",
        "<b>⚠️ The kernel is going to restart. Wait until it's finished before continuing to the next step. ⚠️</b>,\n",
        "</div>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "56e219dbcb9a"
      },
      "source": [
        "### Authenticate your notebook environment (Colab only)\n",
        "Authenticate your environment on Google Colab."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "c97be6a73155"
      },
      "outputs": [],
      "source": [
        "import sys\n",
        "\n",
        "if \"google.colab\" in sys.modules:\n",
        "\n",
        "    from google.colab import auth\n",
        "\n",
        "    auth.authenticate_user()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "before_you_begin:nogpu"
      },
      "source": [
        "## Set Google Cloud project information\n",
        "Learn more about how to [set up a project and a development environment.](https://cloud.google.com/vertex-ai/docs/start/cloud-environment)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "2d96e0c47bed"
      },
      "outputs": [],
      "source": [
        "PROJECT_ID = \"[your-project-id]\"  # @param {type:\"string\"}\n",
        "LOCATION = \"us-central1\"  # @param {type:\"string\"}"
      ]
    },
    {
      "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": [
        "**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 {LOCATION} {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 aiplatform"
      ]
    },
    {
      "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": [
        "aiplatform.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 about [hardware accelerator support](https://cloud.google.com/vertex-ai/docs/general/locations#accelerators) 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 prebuilt containers\n",
        "\n",
        "Set the prebuilt Docker container image for training and prediction.\n",
        "\n",
        "\n",
        "For the latest list, see [prebuilt containers for training](https://cloud.google.com/ai-platform-unified/docs/training/pre-built-containers).\n",
        "\n",
        "\n",
        "For the latest list, see [prebuilt 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 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 don't 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. \n",
        "\n",
        "*Note* : When trainer/task.py is referred to in the worker pool specification, the directory slash is replaced with a dot and the file suffix (.py) is dropped (trainer.task).\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. It's 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 the model artifacts from the command line (--model_dir) are to be saved. If unspecified, then from the environment variable AIP_MODEL_DIR is used.\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, 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. Specify the machine type to provision for training. This informs Vertex AI about the computational resources you need for your training job.\n",
        "  - `machine_type`: The type of GCP instance to provision -- e.g., n1-standard-8.\n",
        "  - `accelerator_type`: The type, if any 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. Configure the disk options , this tells Vertex AI exactly what type and size of disk to allocate for each machine instance during 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 wasn't neccessary to append the `.py` suffix.\n",
        "\n",
        "-`args`: The command line arguments to pass to the corresponding Python module. In this example, you are 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_\" + \"unique\"\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 = aiplatform.CustomJob(\n",
        "    display_name=\"boston_\" + \"unique\", worker_pool_specs=worker_pool_spec\n",
        ")\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 = aiplatform.HyperparameterTuningJob(\n",
        "    display_name=\"boston_\" + \"unique\",\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": [
        "delete_bucket = False\n",
        "\n",
        "# # remove locally generated  directory\n",
        "! rm -r custom\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:\n",
        "# #     ! gsutil rm -r $BUCKET_URI"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "name": "sdk-hyperparameter-tuning.ipynb",
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}