{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1142fd18",
   "metadata": {
    "tags": []
   },
   "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",
   "id": "31a0c126",
   "metadata": {},
   "source": [
    "# Vertex Pipelines: Vertex AI Hyperparameter Tuning Job\n",
    "\n",
    "## Overview\n",
    "This notebook shows how to use the `HyperparameterTuningJobRunOp` to run a hyperparameter tuning job in Vertex AI for a TensorFlow model. While this lab uses TensorFlow for the model code, you could easily replace it with another framework. This sample notebook is based on the [Vertex AI:Hyperparameter Tuning Codelab](https://codelabs.developers.google.com/vertex_hyperparameter_tuning).\n",
    "\n",
    "To learn more about Vertex AI Hyperparameter Tuning Job see [Vertex AI Hyperparameter Tuning Job](https://cloud.google.com/vertex-ai/docs/training/using-hyperparameter-tuning). \n",
    "\n",
    "For `HyperparameterTuningJobRunOp` interface please see the [souce code here](https://github.com/kubeflow/pipelines/blob/master/components/google-cloud/google_cloud_pipeline_components/experimental/hyperparameter_tuning_job)."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "568d5c16",
   "metadata": {},
   "source": [
    "### Install additional packages"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dac98aac",
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip3 install  -U google-cloud-pipeline-components -q\n",
    "!pip3 install  -U google-cloud-aiplatform -q\n",
    "!pip3 install  -U kfp -q"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d0a15440",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Restart the kernel after pip installs\n",
    "import IPython\n",
    "app = IPython.Application.instance()\n",
    "app.kernel.do_shutdown(True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "839e4c48",
   "metadata": {},
   "source": [
    "### Create directory structure"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d1264c91",
   "metadata": {},
   "outputs": [],
   "source": [
    "!mkdir horses_or_humans\n",
    "!mkdir horses_or_humans/trainer"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "907780e7",
   "metadata": {},
   "source": [
    "## Containerize training application code\n",
    "\n",
    "The training application code (inner script) will be put in a Docker container and it will be pushed to the Google Container Registry. After that, the hyperparameter tuning job will be submitted to Vertex by using the `HyperparameterTuningJobRunOp` in a Kubeflow Pipeline. Using this approach, you can tune hyperparameters for a model built with any framework.\n",
    "\n",
    "First, the files below will be created under the a `horses_or_humans` directory. There are several files under that folder:\n",
    "+ Dockerfile\n",
    "+ trainer/\n",
    "    + task.py"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a27d4cee",
   "metadata": {},
   "source": [
    "### Set your Project ID and Pipeline Root"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7e9477a2",
   "metadata": {},
   "outputs": [],
   "source": [
    "PROJECT_ID = \"[your-project-id]\" #@param {type:\"string\"}\n",
    "REGION = \"us-central1\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cb14e1df",
   "metadata": {},
   "source": [
    "### Create a Dockerfile"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "774f7ec3",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%file horses_or_humans/Dockerfile\n",
    "\n",
    "FROM gcr.io/deeplearning-platform-release/tf2-gpu.2-5\n",
    "\n",
    "WORKDIR /\n",
    "\n",
    "# Installs hypertune library\n",
    "RUN pip install cloudml-hypertune\n",
    "\n",
    "# Copies the trainer code to the docker image.\n",
    "COPY trainer /trainer\n",
    "\n",
    "# Sets up the entry point to invoke the trainer.\n",
    "ENTRYPOINT [\"python\", \"-m\", \"trainer.task\"]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8ad8c639",
   "metadata": {},
   "source": [
    "The Dockerfile uses the [Deep Learning Container TensorFlow Enterprise 2.5 GPU Docker image](https://cloud.google.com/ai-platform/deep-learning-containers/docs/choosing-container#choose_a_container_image_type?utm_campaign=CDR_sar_aiml_ucaiplabs_011321&utm_source=external&utm_medium=web). The Deep Learning Containers on Google Cloud come with many common ML and data science frameworks pre-installed. After downloading that image, this Dockerfile sets up the entrypoint for the training code."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bdc2a72a",
   "metadata": {},
   "source": [
    "### Add model training code"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "137a91c8",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%file horses_or_humans/trainer/task.py\n",
    "\n",
    "import tensorflow as tf\n",
    "import tensorflow_datasets as tfds\n",
    "import argparse\n",
    "import hypertune\n",
    "\n",
    "NUM_EPOCHS = 10\n",
    "\n",
    "\n",
    "def get_args():\n",
    "  '''Parses args. Must include all hyperparameters you want to tune.'''\n",
    "\n",
    "  parser = argparse.ArgumentParser()\n",
    "  parser.add_argument(\n",
    "      '--learning_rate',\n",
    "      required=True,\n",
    "      type=float,\n",
    "      help='learning rate')\n",
    "  parser.add_argument(\n",
    "      '--momentum',\n",
    "      required=True,\n",
    "      type=float,\n",
    "      help='SGD momentum value')\n",
    "  parser.add_argument(\n",
    "      '--num_neurons',\n",
    "      required=True,\n",
    "      type=float,\n",
    "      help='number of units in last hidden layer')\n",
    "  args = parser.parse_args()\n",
    "  return args\n",
    "\n",
    "\n",
    "def preprocess_data(image, label):\n",
    "  '''Resizes and scales images.'''\n",
    "\n",
    "  image = tf.image.resize(image, (150,150))\n",
    "  return tf.cast(image, tf.float32) / 255., label\n",
    "\n",
    "\n",
    "def create_dataset():\n",
    "  '''Loads Horses Or Humans dataset and preprocesses data.'''\n",
    "\n",
    "  data, info = tfds.load(name='horses_or_humans', as_supervised=True, with_info=True)\n",
    "\n",
    "  # Create train dataset\n",
    "  train_data = data['train'].map(preprocess_data)\n",
    "  train_data  = train_data.shuffle(1000)\n",
    "  train_data  = train_data.batch(64)\n",
    "\n",
    "  # Create validation dataset\n",
    "  validation_data = data['test'].map(preprocess_data)\n",
    "  validation_data  = validation_data.batch(64)\n",
    "\n",
    "  return train_data, validation_data\n",
    "\n",
    "\n",
    "def create_model(num_neurons, learning_rate, momentum):\n",
    "  '''Defines and complies model.'''\n",
    "\n",
    "  inputs = tf.keras.Input(shape=(150, 150, 3))\n",
    "  x = tf.keras.layers.Conv2D(16, (3, 3), activation='relu')(inputs)\n",
    "  x = tf.keras.layers.MaxPooling2D((2, 2))(x)\n",
    "  x = tf.keras.layers.Conv2D(32, (3, 3), activation='relu')(x)\n",
    "  x = tf.keras.layers.MaxPooling2D((2, 2))(x)\n",
    "  x = tf.keras.layers.Conv2D(64, (3, 3), activation='relu')(x)\n",
    "  x = tf.keras.layers.MaxPooling2D((2, 2))(x)\n",
    "  x = tf.keras.layers.Flatten()(x)\n",
    "  x = tf.keras.layers.Dense(num_neurons, activation='relu')(x)\n",
    "  outputs = tf.keras.layers.Dense(1, activation='sigmoid')(x)\n",
    "  model = tf.keras.Model(inputs, outputs)\n",
    "  model.compile(\n",
    "      loss='binary_crossentropy',\n",
    "      optimizer=tf.keras.optimizers.SGD(learning_rate=learning_rate, momentum=momentum),\n",
    "      metrics=['accuracy'])\n",
    "  return model\n",
    "\n",
    "\n",
    "def main():\n",
    "  args = get_args()\n",
    "  train_data, validation_data = create_dataset()\n",
    "  model = create_model(args.num_neurons, args.learning_rate, args.momentum)\n",
    "  history = model.fit(train_data, epochs=NUM_EPOCHS, validation_data=validation_data)\n",
    "\n",
    "  # DEFINE METRIC\n",
    "  hp_metric = history.history['val_accuracy'][-1]\n",
    "\n",
    "  hpt = hypertune.HyperTune()\n",
    "  hpt.report_hyperparameter_tuning_metric(\n",
    "      hyperparameter_metric_tag='accuracy',\n",
    "      metric_value=hp_metric,\n",
    "      global_step=NUM_EPOCHS)\n",
    "\n",
    "\n",
    "if __name__ == \"__main__\":\n",
    "    main()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "86673e4f",
   "metadata": {},
   "source": [
    "The Python file `task.py` is an inner script that contains the model training code. There are a few components that are specific to using the hyperparameter tuning service.\n",
    "\n",
    "1. The script imports the `hypertune` library. Note that the Dockerfile included instructions to pip install this library.\n",
    "\n",
    "\n",
    "2. The function `get_args()` defines a command-line argument for each hyperparameter you want to tune. In this example, the hyperparameters that will be tuned are the learning rate, the momentum value in the optimizer, and the number of neurons in the last hidden layer of the model. While these are the only hyperparameters targeted here, you are free to modify others. The value passed in those arguments is then used to set the corresponding hyperparameter in the code.\n",
    "\n",
    "3. At the end of the `main()` function, the `hypertune` library is used to define the metric you want to optimize. In TensorFlow, the keras `model.fit` method returns a `History` object. The `History.history` attribute is a record of training loss values and metrics values at successive epochs. If you pass validation data to `model.fit` the `History.history` attribute will include validation loss and metrics values as well. For example, if you trained a model for three epochs with validation data and provided `accuracy` as a metric, the `History.history` attribute would look similar to the following dictionary.\n",
    "```\n",
    "{\n",
    " \"accuracy\": [\n",
    "   0.7795261740684509,\n",
    "   0.9471358060836792,\n",
    "   0.9870933294296265\n",
    " ],\n",
    " \"loss\": [\n",
    "   0.6340447664260864,\n",
    "   0.16712145507335663,\n",
    "   0.04546636343002319\n",
    " ],\n",
    " \"val_accuracy\": [\n",
    "   0.3795261740684509,\n",
    "   0.4471358060836792,\n",
    "   0.4870933294296265\n",
    " ],\n",
    " \"val_loss\": [\n",
    "   2.044623374938965,\n",
    "   4.100203514099121,\n",
    "   3.0728273391723633\n",
    " ]\n",
    "```\n",
    "If you want the hyperparameter tuning service to discover the values that maximize the model's validation accuracy, you can define the metric as the last entry (or `NUM_EPOCHS - 1`) of the `val_accuracy` list. Then, pass this metric to an instance of `HyperTune`. You can pick whatever string you like for the `hyperparameter_metric_tag`, but you’ll need to use the string again later when you kick off the hyperparameter tuning job."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c2df026f",
   "metadata": {},
   "source": [
    "### Build and push the container to the Google Container Registry"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "be3c0914",
   "metadata": {},
   "outputs": [],
   "source": [
    "IMAGE_URI=f\"gcr.io/{PROJECT_ID}/horse-human:hypertune\"\n",
    "%cd horses_or_humans\n",
    "!docker build ./ -t {IMAGE_URI}\n",
    "!docker push {IMAGE_URI}"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f3a09765",
   "metadata": {},
   "source": [
    "## Launch Hyperparameter Tuning Job"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "631a75ba-60d2-4350-b7d1-d88e1756ac56",
   "metadata": {},
   "source": [
    "This section covers launching the Hyperparameter Tuning Job. The syntax uses exact JSON representation of the protos involved, as documented in the [REST API](https://cloud.google.com/vertex-ai/docs/reference/rest/v1/projects.locations.hyperparameterTuningJobs/create) for Vertex AI Hyperparameter Tuning Job. The example here shows how to use utility functions to convert from that of the [HyperparameterTuningJob](https://github.com/googleapis/python-aiplatform/blob/main/google/cloud/aiplatform/jobs.py) in [Vertex AI SDK](https://github.com/googleapis/python-aiplatform) into that of the exact JSON representation, for compatibility."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "89bb4a50",
   "metadata": {},
   "source": [
    "### Import libraries"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0f361e65",
   "metadata": {},
   "outputs": [],
   "source": [
    "from google.cloud.aiplatform import hyperparameter_tuning as hpt\n",
    "from google_cloud_pipeline_components.experimental import hyperparameter_tuning_job\n",
    "from google_cloud_pipeline_components.experimental.custom_job import CustomTrainingJobOp\n",
    "from kfp.v2 import dsl\n",
    "from kfp.v2 import compiler\n",
    "from google.cloud import aiplatform as aip"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d33c87e4-2ada-4b87-bf75-064247f3162d",
   "metadata": {},
   "source": [
    "### Define specs for Hyperparameter Tuning"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "36ceb9f8",
   "metadata": {},
   "outputs": [],
   "source": [
    "# The spec of the worker pools including machine type and Docker image\n",
    "worker_pool_specs = [{\n",
    "    \"machine_spec\": {\n",
    "        \"machine_type\": \"n1-standard-4\",\n",
    "        \"accelerator_type\": \"NVIDIA_TESLA_T4\",\n",
    "        \"accelerator_count\": 1\n",
    "    },\n",
    "    \"replica_count\": 1,\n",
    "    \"container_spec\": {\n",
    "        \"image_uri\": IMAGE_URI\n",
    "    }\n",
    "}]\n",
    "\n",
    "# List serialized from the dictionary representing metrics to optimize.\n",
    "# The dictionary key is the metric_id, which is reported by your training job,\n",
    "# and the dictionary value is the optimization goal of the metric.\n",
    "metric_spec=hyperparameter_tuning_job.serialize_metrics({'accuracy': 'maximize'})\n",
    "\n",
    "# List serialized from the parameter dictionary. The dictionary\n",
    "# represents parameters to optimize. The dictionary key is the parameter_id,\n",
    "# which is passed into your training job as a command line key word argument, and the\n",
    "# dictionary value is the parameter specification of the metric.\n",
    "parameter_spec = hyperparameter_tuning_job.serialize_parameters({\n",
    "    \"learning_rate\": hpt.DoubleParameterSpec(min=0.001, max=1, scale=\"log\"),\n",
    "    \"momentum\": hpt.DoubleParameterSpec(min=0, max=1, scale=\"linear\"),\n",
    "    \"num_neurons\": hpt.DiscreteParameterSpec(values=[64, 128, 512], scale=None)\n",
    "})"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d92d9ae7",
   "metadata": {},
   "source": [
    "### Define the pipeline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2a67cde8",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "PIPELINE_ROOT = 'gs://[your-base-output-directory]'  #@param {type:\"string\"}\n",
    "    \n",
    "@dsl.pipeline(pipeline_root=PIPELINE_ROOT, name='hp-tune-pipeline')\n",
    "def hp_tune_pipeline():\n",
    "\n",
    "    hp_tuning_task = hyperparameter_tuning_job.HyperparameterTuningJobRunOp(\n",
    "        display_name='hp-job',\n",
    "        project=PROJECT_ID,\n",
    "        location=REGION,\n",
    "        worker_pool_specs=worker_pool_specs,\n",
    "        study_spec_metrics=metric_spec,\n",
    "        study_spec_parameters=parameter_spec,\n",
    "        max_trial_count=15,\n",
    "        parallel_trial_count=3,\n",
    "        base_output_directory=PIPELINE_ROOT\n",
    "    )\n",
    "    \n",
    "    trials_task = hyperparameter_tuning_job.GetTrialsOp(\n",
    "      gcp_resources=hp_tuning_task.outputs['gcp_resources'])\n",
    "\n",
    "    best_trial_task = hyperparameter_tuning_job.GetBestTrialOp(\n",
    "      trials=trials_task.output, study_spec_metrics=metric_spec)\n",
    "\n",
    "    is_accuracy_beyond_threshold_task = hyperparameter_tuning_job.IsMetricBeyondThresholdOp(\n",
    "      trial=best_trial_task.output, study_spec_metrics=metric_spec, threshold=0.7)\n",
    "\n",
    "    with dsl.Condition(\n",
    "        is_accuracy_beyond_threshold_task.output == \"true\",\n",
    "        name=\"deploy_decision\",\n",
    "    ):\n",
    "        best_hyperparameters_task = hyperparameter_tuning_job.GetHyperparametersOp(\n",
    "          trial=best_trial_task.output)\n",
    "\n",
    "        # Construct new worker_pool_specs based on best hyperparameters\n",
    "        worker_pool_specs_task = hyperparameter_tuning_job.GetWorkerPoolSpecsOp(\n",
    "          best_hyperparameters=best_hyperparameters_task.output,\n",
    "          worker_pool_specs=worker_pool_specs\n",
    "        )\n",
    "\n",
    "        # Train new model based on new worker_pool_specs\n",
    "        training_task = CustomTrainingJobOp(\n",
    "          project=PROJECT_ID,\n",
    "          location=REGION,\n",
    "          display_name='training-job',\n",
    "          worker_pool_specs=worker_pool_specs_task.output\n",
    "        )"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3211ba19",
   "metadata": {
    "tags": []
   },
   "source": [
    "### Compile and run the pipeline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c368c73f",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "job = aip.PipelineJob.from_pipeline_func(\n",
    "    pipeline_func=hp_tune_pipeline, location = \"us-west1\", enable_caching=False\n",
    ")\n",
    "\n",
    "job.run()"
   ]
  }
 ],
 "metadata": {
  "environment": {
   "kernel": "python3",
   "name": "common-cpu.m86",
   "type": "gcloud",
   "uri": "gcr.io/deeplearning-platform-release/base-cpu:m86"
  },
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
