notebooks/official/tabnet/tabnet_vertex_tutorial.ipynb

{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ur8xi4C7S06n"
      },
      "outputs": [],
      "source": [
        "# Copyright 2022 Google LLC\n",
        "#\n",
        "# Licensed under the Apache License, Version 2.0 (the \"License\");\n",
        "# you may not use this file except in compliance with the License.\n",
        "# You may obtain a copy of the License at\n",
        "#\n",
        "#     https://www.apache.org/licenses/LICENSE-2.0\n",
        "#\n",
        "# Unless required by applicable law or agreed to in writing, software\n",
        "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
        "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
        "# See the License for the specific language governing permissions and\n",
        "# limitations under the License."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "JAPoU8Sm5E6e"
      },
      "source": [
        "# Vertex AI TabNet\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/tabnet/tabnet_vertex_tutorial.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/tabnet/tabnet_vertex_tutorial.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/tabnet/tabnet_vertex_tutorial.ipynb\">\n",
        "      <img src=\"https://lh3.googleusercontent.com/UiNooY4LUgW_oTvpsNhPpQzsstV5W8F7rYgxgGBD85cWJoLmrOzhVs_ksK_vgx40SHs7jCqkTkCk=e14-rj-sc0xffffff-h130-w32\" alt=\"Vertex AI logo\">\n",
        "      Open in Vertex AI Workbench\n",
        "    </a>\n",
        "  </td>                                                                                               \n",
        "</table>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "2a0f26a1bc15"
      },
      "source": [
        "## Overview\n",
        "\n",
        "This tutorial demonstrates how to use the TabNet algorithm on the Vertex AI platform.\n",
        "\n",
        "TabNet combines the best of two worlds: it is explainable (similar to simpler tree-based models) while benefiting from high performance (similar to deep neural networks). This makes it great for retailers, finance and insurance industry applications such as predicting credit scores, fraud detection and forecasting. \n",
        "\n",
        "TabNet uses a machine learning technique called sequential attention to select which model features to reason from at each step in the model. This mechanism makes it possible to explain how the model arrives at its predictions and helps it learn more accurate models. Thanks to this design, TabNet not only outperforms other neural networks and decision trees but also provides interpretable feature attributions. Releasing TabNet as a First Party Trainer in Vertex AI means you'll be able to easily take advantage of TabNet's architecture and explainability and use it to train models on your own data. \n",
        "\n",
        "Learn more about [Tabular Workflow for TabNet](https://cloud.google.com/vertex-ai/docs/tabular-data/tabular-workflows/tabnet)."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "tvgnzT1CKxrO"
      },
      "source": [
        "### Objective\n",
        "\n",
        "In this notebook, you learn how to run TabNet model on Vertex AI.\n",
        "\n",
        "This tutorial uses the following Google Cloud ML services and resources:\n",
        "\n",
        "- Vertex AI Training\n",
        "- Vertex AI Hyperparameter Tuning\n",
        "- TabNet builtin algorithm\n",
        "- BigQuery\n",
        "\n",
        "The steps performed are:\n",
        "1. **Setup**: Importing the required libraries and setting your global variables.\n",
        "2. **Configure parameters**: Setting the appropriate parameter values for the training job.\n",
        "3. **Train on Vertex AI Training**: Submitting a training job using csv input.\n",
        "4. **Hyperparameter tuning**: Running a hyperparameter tuning job.\n",
        "5. **Hyperparameter on Vertex AI Training with BigQuery input**: Submitting a training job using BigQuery input.\n",
        "6. **Cleaning up**: Deleting resources created by this tutorial."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "d60d50a79198"
      },
      "source": [
        "### Dataset\n",
        "\n",
        "This tutorial uses the `petfinder` in the public bucket `gs://cloud-samples-data/ai-platform-unified/datasets/tabular/`, which was generated from the [PetFinder.my Adoption Prediction](https://www.kaggle.com/c/petfinder-adoption-prediction). "
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "0c997d8d92ce"
      },
      "source": [
        "### Costs \n",
        "\n",
        "\n",
        "This tutorial uses billable components of Google Cloud:\n",
        "\n",
        "* Vertex AI\n",
        "* Cloud Storage\n",
        "\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": "xOMNWzTbftDr"
      },
      "source": [
        "## Installation\n",
        "\n",
        "Install the following packages required to execute this notebook. "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "wyy5Lbnzg5fi"
      },
      "outputs": [],
      "source": [
        "! pip3 install --quiet --upgrade tensorflow\n",
        "! pip3 install --quiet --upgrade google-cloud-aiplatform \\\n",
        "                                 tensorboard-plugin-profile \\\n",
        "                                 google-cloud-pipeline-components\n",
        "! gcloud components update --quiet"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "restart"
      },
      "source": [
        "### Colab only: Uncomment the following cell to restart the kernel.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "bzPxhxS5lugp"
      },
      "outputs": [],
      "source": [
        "# Automatically restart kernel after installs so that your environment can access the new packages\n",
        "# import IPython\n",
        "\n",
        "# app = IPython.Application.instance()\n",
        "# app.kernel.do_shutdown(True)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "d2qpIurSjmpT"
      },
      "source": [
        "## Before you begin\n",
        "\n",
        "### Set up your Google Cloud project\n",
        "\n",
        "**The following steps are required, regardless of your notebook environment.**\n",
        "\n",
        "1. [Select or create a Google Cloud project](https://console.cloud.google.com/cloud-resource-manager). When you first create an account, you get a $300 free credit towards your compute/storage costs.\n",
        "\n",
        "2. [Make sure that billing is enabled for your project](https://cloud.google.com/billing/docs/how-to/modify-project).\n",
        "\n",
        "3. [Enable the following APIs: Vertex AI API, Cloud Resource Manager API](https://console.cloud.google.com/flows/enableapi?apiid=aiplatform.googleapis.com,cloudresourcemanager.googleapis.com).\n",
        "\n",
        "4. If you are running this notebook locally, you need to install the [Cloud SDK](https://cloud.google.com/sdk)."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "project_id"
      },
      "source": [
        "#### Set your project ID\n",
        "\n",
        "**If you don't know your project ID**, try the following:\n",
        "* Run `gcloud config list`.\n",
        "* Run `gcloud projects list`.\n",
        "* See the support page: [Locate the project ID](https://support.google.com/googleapi/answer/7014113)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "wsePm9c4jmpT"
      },
      "outputs": [],
      "source": [
        "PROJECT_ID = \"[your-project-id]\"  # @param {type:\"string\"}\n",
        "\n",
        "# Set the project id\n",
        "! gcloud config set project {PROJECT_ID}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "a54f9d7c1876"
      },
      "source": [
        "#### Region\n",
        "\n",
        "You can also change the `REGION` variable used by Vertex AI. Learn more about [Vertex AI regions](https://cloud.google.com/vertex-ai/docs/general/locations)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "3aaadaaf9b30"
      },
      "outputs": [],
      "source": [
        "REGION = \"us-central1\"  # @param {type: \"string\"}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "5c0404984792"
      },
      "source": [
        "### Authenticate your Google Cloud account\n",
        "\n",
        "Depending on your Jupyter environment, you may have to manually authenticate. Follow the relevant instructions below.\n",
        "\n",
        "**1. Vertex AI Workbench**\n",
        "* Do nothing as you are already authenticated.\n",
        "\n",
        "**2. Local JupyterLab instance, uncomment and run:**"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Nt8cEM2GjmpU"
      },
      "outputs": [],
      "source": [
        "# ! gcloud auth login"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "XUSL_JcpjmpU"
      },
      "source": [
        "**3. Colab, uncomment and run:**"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "_2zemfGvjmpU"
      },
      "outputs": [],
      "source": [
        "# from google.colab import auth\n",
        "# auth.authenticate_user()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "TCPJ38n7jmpU"
      },
      "source": [
        "**4. Service account or other**\n",
        "* See how to grant Cloud Storage permissions to your service account at https://cloud.google.com/storage/docs/gsutil/commands/iam#ch-examples."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "bucket:custom"
      },
      "source": [
        "### Create a Cloud Storage bucket\n",
        "\n",
        "Create a storage bucket to store intermediate artifacts such as datasets.\n",
        "\n",
        "When you submit a training job using the Cloud SDK, you upload a Python package\n",
        "containing your training code to a Cloud Storage bucket. Vertex AI runs\n",
        "the code from this package. In this tutorial, Vertex AI also saves the\n",
        "trained model that results from your job in the same bucket. Using this model artifact, you can then\n",
        "create Vertex AI Model resource and use for prediction."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "bucket"
      },
      "outputs": [],
      "source": [
        "BUCKET_URI = f\"gs://your-bucket-name-{PROJECT_ID}-unique\"  # @param {type:\"string\"}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "create_bucket"
      },
      "source": [
        "**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": "Oz8J0vmSlugt"
      },
      "outputs": [],
      "source": [
        "! gsutil mb -l {REGION} -p {PROJECT_ID} {BUCKET_URI}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "XoEqT2Y4DJmf"
      },
      "source": [
        "### Import libraries and define constants"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "pRUOFELefqf1"
      },
      "outputs": [],
      "source": [
        "import os\n",
        "import re\n",
        "import time\n",
        "from datetime import datetime\n",
        "\n",
        "from google.cloud import aiplatform\n",
        "\n",
        "%load_ext tensorboard"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "BoVqAMOuedPe"
      },
      "outputs": [],
      "source": [
        "# Please note that if you use csv input, the first column is the label column.\n",
        "\n",
        "IMPORT_FILE = \"petfinder-tabular-classification-tabnet-with-header.csv\"\n",
        "! gsutil cp gs://cloud-samples-data/ai-platform-unified/datasets/tabular/{IMPORT_FILE} {BUCKET_URI}/data/petfinder/train.csv\n",
        "\n",
        "gcs_source = f\"{BUCKET_URI[5:]}/data/petfinder/train.csv\"\n",
        "print(\"gcs_source: \", gcs_source)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "CMCfG2ao8kLh"
      },
      "source": [
        "## Configure parameters"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "hxB0FTQT8uqY"
      },
      "source": [
        "The following table shows parameters that are common to all Vertex Training jobs created using the `gcloud ai custom-jobs create` command. See the [official documentation](https://cloud.google.com/sdk/gcloud/reference/ai/custom-jobs/create) for all the possible arguments.\n",
        "\n",
        "| Parameter | Data type | Description | Required |\n",
        "|--|--|--|--|\n",
        "| `display-name` | string | Name of the job. | Yes |\n",
        "| `worker-pool-spec` | string | Comma-separated list of arguments specifying a worker pool configuration (see below). | Yes |\n",
        "| `region` | string | Region to submit the job to. | No |\n",
        "\n",
        "The `worker-pool-spec` flag can be specified multiple times, one for each worker pool. The following table shows the arguments used to specify a worker pool.\n",
        "\n",
        "| Parameter | Data type | Description | Required |\n",
        "|--|--|--|--|\n",
        "| `machine-type` | string | Machine type for the pool. See the [official documentation](https://cloud.google.com/vertex-ai/docs/training/configure-compute) for supported machines. | Yes |\n",
        "| `replica-count` | int | The number of replicas of the machine in the pool. | No |\n",
        "| `container-image-uri` | string | Docker image to run on each worker. | No |"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Xc_I_XHz8z6B"
      },
      "source": [
        "The following table shows the parameters for the TabNet training job:\n",
        "\n",
        "| Parameter | Data type | Description | Required |\n",
        "|--|--|--|--|\n",
        "| `preprocess` | boolean argument | Specify this to enable automatic preprocessing. | No|\n",
        "| `job_dir` | string | Cloud Storage directory where the model output files will be stored. | Yes |\n",
        "| `input_metadata_path` | string | The GCS path to the TabNet-specific metadata for the training dataset. Please see above on how to create the metadata. | No . |\n",
        "| `training_data_path` | string | Cloud Storage pattern where training data is stored. | Yes |\n",
        "| `validation_data_path` | string | Cloud Storage pattern where eval data is stored. | No |\n",
        "| `test_data_path` | string | Cloud Storage pattern where test data is stored. | Yes |\n",
        "| `input_type` | string | “bigquery“ or “csv“ - type of the input tabular data. If csv is mentioned then the first column is treated as target. If CSV files have a header, also pass the flag “data_has_header”. If “bigquery” is used, one can either supply training/validation data paths, or supply BigQuery project, dataset, and table names for preprocessing to produce the training and validation datasets.. | No -Deafult is \"csv\" |\n",
        "| `model_type` | string | The learning task such as classification or regression. | Yes |\n",
        "| `split_column` | string | The column name used to create the training, validation, and test splits. Values of the columns (a.k.a table['split_column']) should be either “TRAIN”, “VALIDATE”, or “TEST”. “TEST” is optional. Applicable to bigquery input only. | No. |\n",
        "| `train_batch_size` | int | Batch size for training. | No - Default is 1024. |\n",
        "| `eval_split` | float | Split fraction to use for the evaluation dataset, if `validation_data_path` is not provided. | No - Default is 0.2 |\n",
        "| `learning_rate` | float | Learning rate for training. | No - Default is the default learning rate of the specified optimizer. |\n",
        "| `eval_frequency_secs` | int | Frequency at which evaluation and checkpointing will take place.The default is 600. | No . |\n",
        "| `num_parallel_reads` | int | Number of threads used to read input files. We suggest setting it equal or slightly less than the number of CPUs of the machine for maximal performance in most cases. For example, 6 per GPU is a good default choice. | Yes . |\n",
        "| `data_cache` | string | Choose to cache data to “memory”, “disk” or “no_cache”. For large datasets, caching the data into memory would throw out-of-memory errors, therefore, we suggest choosing “disk”. You can specify the disk size in the config file (as exemplified below). Make sure to request a sufficiently large (e.g. TB size) disk to write the data, for large (B-scale) datasets. | No, The default one is “memory” . |\n",
        "| `bq_project` | string | The name of the BigQuery project. If input_type=bigquery and using the flag –preprocessing, this is required. This is an alternative to specifying train, validation, and test data paths. | No . |\n",
        "| `dataset_name` | string | The name of the BigQuery dataset. If input_type=bigquery and using the flag –preprocessing, this is required. This is an alternative to specifying train, validation, and test data paths. | No . |\n",
        "| `table_name` | string | The name of the BigQuery table. If input_type=bigquery and using the flag –preprocessing, this is required. This is an alternative to specifying train, validation, and test data paths. | No . |\n",
        "| `loss_function_type` | string | There are several loss function types in TabNet. For regression: mse/mae are included. For classification: cross_entropy/weighted_cross_entropy/focal_loss are included | No . If the value is \"default\", we use mse for regression and cross_entropy for classification.|\n",
        "| `deterministic_data` | boolean argument | Determinism of data reading from tabular data. The default is set to False. When setting to True, the experiment is deterministic. For fast training on large datasets, we suggest deterministic_data=False setting, albeit having randomness in the results (which becomes negligible at large datasets). Note that determinism is still not guaranteed with distributed training, as map-reduce causes randomness due to the ordering of algebraic operations with finite precision. However, this is negligible in practice, especially on large datasets. For the cases where 100% determinism is desired, besides deterministic_data=True setting, we suggest training with a single GPU (e.g. with MACHINE_TYPE=\"n1-highmem-8\").| No, The default is False . |\n",
        "| `stream_inputs` | boolean argument | Stream input data from GCS instead of downloading it locally - this option is suggested for fast runtime. | No . |\n",
        "| `large_category_dim` | int | Dimensionality of the embedding - if the number of distinct categories for a categorical column is greater than large_category_thresh we use a large_category_dim dimensional embedding, instead of 1-D embedding. The default is 1. We suggest increasing it (e.g. to ~5 in most cases and even ~10 if the number of categories is typically very large in the dataset), if pushing the accuracy is the main goal, rather than computational efficiency and explainability.  | No . |\n",
        "| `large_category_thresh` | int | Threshold for categorical column cardinality - if the number of distinct categories for a categorical column is greater than large_category_thresh we use a large_category_dim dimensional embedding, instead of 1-D embedding. The default is 300. We suggest decreasing it (e.g. to ~10), if pushing the accuracy is the main goal, rather than computational efficiency and explainability. | No . |\n",
        "| `yeo_johnson_transform` | boolean argument | Enables trainable Yeo-Johnson Power Transform (the default is disabled). Please see this link: https://www.stat.umn.edu/arc/yjpower.pdf for more information on Yeo-Johnson Power Transform. With our implementation, the transform parameters are learnable along with TabNet, trained in an end-to-end way. | No . |\n",
        "| `apply_log_transform` | boolean argument | If log transform statistics are contained in the metadata and this flag is true then the input features will be log transformed. Use false to not use the transform and true (the default) to use it. Especially for datasets with skewed numerical distributions, log transformations could be very helpful.  | No . |\n",
        "| `apply_quantile_transform` | boolean argument | If quantile statistics are contained in the metadata and this flag is true then the input features will be quantile transformed. Use false to not use the transform and true (the default) to use it. Especially for datasets with skewed numerical distributions, quantile transformations could be very helpful. Currently supported for BigQuery input_type. | No . |\n",
        "| `replace_transformed_features` | boolean argument | If true then if a transformation is applied to a feature that feature will be replaced. If false (the default option) then the transformed feature 'will be added as a new feature to the list of feature columns. Use true to replace the feature with the transform and false to append the transformed feature as a new column. | No . |\n",
        "| `target_column` | string | name of the label column. Note that for classification, the label needs to be of String or integer type. | No . |\n",
        "| `prediction_raw_inputs` | boolean argument | If we set this argument, the model serving allows us to pass the features as a dictionary of tensors instead of CSV row. | No . |\n",
        "| `exclude_key` | boolean argument | If we set this argument, we exclude a key in the input/output. The key is helpful in batch prediction processes input and saves output in an unpredictable order. The key helps match the output with input. | No . |"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "FNmQS6EW86gp"
      },
      "outputs": [],
      "source": [
        "# URI of the TabNet training Docker image.\n",
        "LEARNER_IMAGE_URI = (\n",
        "    \"us-docker.pkg.dev/vertex-ai-restricted/builtin-algorithm/tab_net_v2\"\n",
        ")\n",
        "\n",
        "DATASET_NAME = \"petfinder\"  # Change to your dataset name.\n",
        "\n",
        "ALGORITHM = \"tabnet\"\n",
        "MODEL_TYPE = \"classification\"\n",
        "MODEL_NAME = f\"{DATASET_NAME}_{ALGORITHM}_{MODEL_TYPE}\"\n",
        "\n",
        "TRAINING_DATA_PATH = f\"{BUCKET_URI}/data/{DATASET_NAME}/train.csv\"\n",
        "\n",
        "print(\"training path: \", TRAINING_DATA_PATH)\n",
        "print(\"model name: \", MODEL_NAME)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "k1gLmZCPnm_s"
      },
      "source": [
        "# Run TabNet training"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "7W2rij5opa8R"
      },
      "outputs": [],
      "source": [
        "current_time = datetime.utcnow().strftime('%y%m%d%H%M%S')\n",
        "OUTPUT_DIR = f\"{BUCKET_URI}/{MODEL_NAME}_{current_time}\"\n",
        "print('Output dir: ', OUTPUT_DIR)\n",
        "learning_job_name = f'tab_net_{DATASET_NAME}_{current_time}'\n",
        "print('Job name: ', learning_job_name)\n",
        "\n",
        "CREATION_LOG = !gcloud ai custom-jobs create \\\n",
        "  --region={REGION} \\\n",
        "  --display-name={learning_job_name} \\\n",
        "  --worker-pool-spec=machine-type=n1-standard-8,replica-count=1,container-image-uri={LEARNER_IMAGE_URI} \\\n",
        "  --args=--preprocess \\\n",
        "  --args=--prediction_raw_inputs \\\n",
        "  --args=--exclude_key \\\n",
        "  --args=--data_has_header \\\n",
        "  --args=--training_data_path={TRAINING_DATA_PATH} \\\n",
        "  --args=--job-dir={OUTPUT_DIR} \\\n",
        "  --args=--model_type={MODEL_TYPE} \\\n",
        "  --args=--max_steps=2000 \\\n",
        "  --args=--batch_size=4096 \\\n",
        "  --args=--learning_rate=0.01\n",
        "\n",
        "print(CREATION_LOG)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "y_GgAlhpGl4q"
      },
      "source": [
        "After the job is submitted successfully, you can view its details and logs:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Lh1ZQ9uRGdk-"
      },
      "outputs": [],
      "source": [
        "JOB_ID = re.search(r\"(?<=/customJobs/)\\d+\", CREATION_LOG[1]).group(0)\n",
        "print(JOB_ID)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "IShlAxz8piAg"
      },
      "outputs": [],
      "source": [
        "# View the job's configuration and state.\n",
        "STATE = \"state: JOB_STATE_PENDING\"\n",
        "\n",
        "while STATE not in [\"state: JOB_STATE_SUCCEEDED\", \"state: JOB_STATE_FAILED\"]:\n",
        "    DESCRIPTION = ! gcloud ai custom-jobs describe {JOB_ID} --region={REGION}\n",
        "    STATE = DESCRIPTION[-2]\n",
        "    print(STATE)\n",
        "    time.sleep(60)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "2ed50da51105"
      },
      "source": [
        "# Deploy on Vertex AI Prediction\n",
        "\n",
        "## Import the model\n",
        "Your training job exports two TF SavedModel under `gs://<job_dir>/model`. The exported model can be used for online or batch prediction in Vertex AI Prediction. First, import the model:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "e9dee2e91da8"
      },
      "outputs": [],
      "source": [
        "# Imports the model.\n",
        "current_time = datetime.utcnow().strftime(\"%y%m%d%H%M%S\")\n",
        "MODEL_TYPE = \"classification\"\n",
        "\n",
        "DISPLAY_NAME = f\"tabnet_{DATASET_NAME}_{MODEL_TYPE}_{current_time}\"  # The display name of the model.\n",
        "MODEL_NAME = f\"tabnet_{MODEL_TYPE}_model\"  # Used by the deployment container.\n",
        "ARTIFACT_URI = f\"{OUTPUT_DIR}/model/\"\n",
        "\n",
        "print(\"Display_name:\", DISPLAY_NAME)\n",
        "print(\"Model name: \", MODEL_NAME)\n",
        "\n",
        "aiplatform.init(\n",
        "    project=PROJECT_ID,\n",
        "    location=REGION,\n",
        "    staging_bucket=BUCKET_URI,\n",
        ")\n",
        "\n",
        "model = aiplatform.Model.upload(\n",
        "    display_name=DISPLAY_NAME,\n",
        "    artifact_uri=ARTIFACT_URI,\n",
        "    serving_container_image_uri=\"us-docker.pkg.dev/vertex-ai/prediction/tf2-cpu.2-4:latest\",\n",
        ")\n",
        "print(model)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "6aad41003302"
      },
      "source": [
        "## Deploy the model\n",
        "After importing the model, you must deploy it to an endpoint so that you can get online predictions. More information about this process can be found in the [official documentation](https://cloud.google.com/vertex-ai/docs/predictions/deploy-model-api)."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "d463a410ac44"
      },
      "source": [
        "## Create a model endpoint"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "10224e3190c3"
      },
      "outputs": [],
      "source": [
        "endpoint = aiplatform.Endpoint.create(display_name=DATASET_NAME)\n",
        "\n",
        "# If you want to use existing endpoint\n",
        "# endpoint = aiplatform.Endpoint('[location of existing endpoint]')\n",
        "print(\"Endpoint: \", endpoint)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "33fe5010c84e"
      },
      "source": [
        "## Deploy model to the endpoint"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "d29efd150e44"
      },
      "outputs": [],
      "source": [
        "model.deploy(\n",
        "    endpoint=endpoint,\n",
        "    machine_type=\"n1-standard-4\",\n",
        "    traffic_split={\"0\": 100},\n",
        "    deployed_model_display_name=DISPLAY_NAME,\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "92aaeb09cd63"
      },
      "source": [
        "## Online prediction"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "2288d32ddfb8"
      },
      "outputs": [],
      "source": [
        "prediction = endpoint.predict(\n",
        "    [\n",
        "        {\n",
        "            \"Type\": \"Cat\",\n",
        "            \"Age\": 3,\n",
        "            \"Breed1\": \"Tabby\",\n",
        "            \"Gender\": \"Male\",\n",
        "            \"Color1\": \"Black\",\n",
        "            \"Color2\": \"White\",\n",
        "            \"MaturitySize\": \"Small\",\n",
        "            \"FurLength\": \"Short\",\n",
        "            \"Vaccinated\": \"No\",\n",
        "            \"Sterilized\": \"No\",\n",
        "            \"Health\": \"Healthy\",\n",
        "            \"Fee\": 100,\n",
        "            \"PhotoAmt\": 2,\n",
        "        }\n",
        "    ]\n",
        ")\n",
        "\n",
        "print(prediction)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "E9EZvfSUWrxS"
      },
      "source": [
        "# Hyperparameter Tuning\n",
        "After successfully training your model, deploying it, and calling it to make predictions, you may want to optimize the hyperparameters used during training to improve your model's accuracy and performance. See the Vertex AI documentation for an overview of hyperparameter tuning and how to use it in your Vertex Training jobs.\n",
        "\n",
        "For this example, the following command runs a Vertex AI hyperparameter tuning job with 4 trials that attempts to maximize the validation AUC metric. The hyperparameters it optimizes are the number of max_steps and the learning rate."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "DnJcD4vBeKtl"
      },
      "outputs": [],
      "source": [
        "DATASET_NAME = \"petfinder\"  # Change to your dataset name.\n",
        "current_time = datetime.utcnow().strftime('%y%m%d%H%M%S')\n",
        "OUTPUT_DIR = f\"{BUCKET_URI[5:]}/hptuning/{MODEL_NAME}_{current_time}/\"\n",
        "print('Output dir: ', OUTPUT_DIR)\n",
        "learning_job_name = f'tab_net_{DATASET_NAME}_{current_time}'\n",
        "print('Job name: ', learning_job_name)\n",
        "\n",
        "config = f\"\"\"studySpec:\n",
        "  metrics:\n",
        "  - metricId: auc\n",
        "    goal: MAXIMIZE\n",
        "  parameters:\n",
        "  - parameterId: max_steps\n",
        "    integerValueSpec:\n",
        "      minValue: 2000\n",
        "      maxValue: 3000\n",
        "  - parameterId: learning_rate\n",
        "    doubleValueSpec:\n",
        "      minValue: 0.0000001\n",
        "      maxValue: 0.1\n",
        "trialJobSpec:\n",
        "  workerPoolSpecs:\n",
        "  - machineSpec:\n",
        "      machineType: n1-highmem-8\n",
        "      acceleratorType: NVIDIA_TESLA_V100\n",
        "      acceleratorCount: 1\n",
        "    replicaCount: 1\n",
        "    diskSpec:\n",
        "      bootDiskType: pd-ssd\n",
        "      bootDiskSizeGb: 100\n",
        "    containerSpec:\n",
        "      imageUri: {LEARNER_IMAGE_URI}\n",
        "      args:\n",
        "      - --preprocess \n",
        "      - --prediction_raw_inputs\n",
        "      - --exclude_key\n",
        "      - --data_has_header\n",
        "      - --training_data_path={TRAINING_DATA_PATH}\n",
        "      - --job-dir={OUTPUT_DIR}\n",
        "      - --batch_size=1028\n",
        "      - --model_type=classification\n",
        "\"\"\"\n",
        "\n",
        "!echo $'{config}' > ./config.yaml\n",
        "\n",
        "MAX_TRIAL_COUNT=4\n",
        "PARALLEL_TRIAL_COUNT=2\n",
        "\n",
        "!gcloud ai hp-tuning-jobs create \\\n",
        "  --config=config.yaml \\\n",
        "  --max-trial-count={MAX_TRIAL_COUNT} \\\n",
        "  --parallel-trial-count={PARALLEL_TRIAL_COUNT} \\\n",
        "  --region=$REGION \\\n",
        "  --display-name=$learning_job_name"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "mpKDxrP9yQK8"
      },
      "source": [
        "# Hyperparameter Tuning with Big Query (BQ) input\n",
        "You need to create the BQ dataset using the above csv file."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "2z6gsQQiVQq_"
      },
      "source": [
        "Create BQ dataset to your project\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "NXu00dB7VZ8w"
      },
      "outputs": [],
      "source": [
        "! bq --location={REGION} mk --dataset {PROJECT_ID}:{DATASET_NAME}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "0JfQ_Lz_DHUP"
      },
      "source": [
        "Create the BQ table using the csv file."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "zv5ZK67LXh_e"
      },
      "outputs": [],
      "source": [
        "!bq --location={REGION} load --source_format=CSV --autodetect {PROJECT_ID}:{DATASET_NAME}.train {BUCKET_URI}/data/petfinder/train.csv"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Qsz_ACeWVPb0"
      },
      "source": [
        "Run hp-tuning using BQ input."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "pPaqOq6_UGIl"
      },
      "outputs": [],
      "source": [
        "current_time = datetime.utcnow().strftime('%y%m%d%H%M%S')\n",
        "\n",
        "BQ_PROJECT=f\"{PROJECT_ID}\"\n",
        "DATASET_NAME = \"petfinder\"  # Change to your dataset name.\n",
        "TABLE_NAME=\"train\"\n",
        "OUTPUT_DIR = f\"{BUCKET_URI}/hptuning/{MODEL_NAME}_{current_time}/\"\n",
        "print('Output dir: ', OUTPUT_DIR)\n",
        "learning_job_name = f'tab_net_{DATASET_NAME}_{current_time}'\n",
        "print('Job name: ', learning_job_name)\n",
        "\n",
        "config = f\"\"\"studySpec:\n",
        "  metrics:\n",
        "  - metricId: auc\n",
        "    goal: MAXIMIZE\n",
        "  parameters:\n",
        "  - parameterId: max_steps\n",
        "    integerValueSpec:\n",
        "      minValue: 2000\n",
        "      maxValue: 3000\n",
        "  - parameterId: learning_rate\n",
        "    doubleValueSpec:\n",
        "      minValue: 0.0000001\n",
        "      maxValue: 0.1\n",
        "trialJobSpec:\n",
        "  workerPoolSpecs:\n",
        "  - machineSpec:\n",
        "      machineType: n1-highmem-8\n",
        "      acceleratorType: NVIDIA_TESLA_V100\n",
        "      acceleratorCount: 1\n",
        "    replicaCount: 1\n",
        "    diskSpec:\n",
        "      bootDiskType: pd-ssd\n",
        "      bootDiskSizeGb: 200\n",
        "    containerSpec:\n",
        "      imageUri: {LEARNER_IMAGE_URI}\n",
        "      args:\n",
        "      - --eval_frequency_secs=10800\n",
        "      - --input_type=bigquery\n",
        "      - --preprocess\n",
        "      - --prediction_raw_inputs\n",
        "      - --exclude_key\n",
        "      - --model_type=classification\n",
        "      - --stream_inputs\n",
        "      - --bq_project={BQ_PROJECT}\n",
        "      - --dataset_name={DATASET_NAME}\n",
        "      - --table_name={TABLE_NAME}\n",
        "      - --target_column=Adopted\n",
        "      - --num_parallel_reads=2\n",
        "      - --optimizer_type=adam\n",
        "      - --data_cache=disk\n",
        "      - --deterministic_data=False\n",
        "      - --loss_function_type=weighted_cross_entropy\n",
        "      - --replace_transformed_features=True\n",
        "      - --apply_quantile_transform=True\n",
        "      - --apply_log_transform=True\n",
        "      - --batch_size=32768\n",
        "      - --job-dir={OUTPUT_DIR}\n",
        "\"\"\"\n",
        "\n",
        "!echo $'{config}' > ./hptuning-config.yaml\n",
        "\n",
        "MAX_TRIAL_COUNT=2\n",
        "PARALLEL_TRIAL_COUNT=2\n",
        "\n",
        "!gcloud ai hp-tuning-jobs create \\\n",
        "  --config=hptuning-config.yaml \\\n",
        "  --max-trial-count={MAX_TRIAL_COUNT} \\\n",
        "  --parallel-trial-count={PARALLEL_TRIAL_COUNT} \\\n",
        "  --region=$REGION \\\n",
        "  --display-name=$learning_job_name"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "TpV-iwP9qw9c"
      },
      "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": "sx_vKniMq9ZX"
      },
      "outputs": [],
      "source": [
        "# Delete Cloud Storage objects that were created\n",
        "! gsutil -m rm -r $JOB_DIR\n",
        "\n",
        "# Delete BQ table\n",
        "! bq rm -f {PROJECT_ID}:{DATASET_NAME}.train\n",
        "\n",
        "# Delete endpoint resource\n",
        "endpoint.delete(force=True)\n",
        "\n",
        "# Delete model resource\n",
        "model.delete()\n",
        "\n",
        "if os.getenv(\"IS_TESTING\"):\n",
        "    ! gsutil -m rm -r $BUCKET_URI"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "collapsed_sections": [],
      "name": "tabnet_vertex_tutorial.ipynb",
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}