notebooks/vertex_nas_pointpillars_tfvision.ipynb

{
  "cells": [
   {
    "cell_type": "markdown",
    "metadata": {},
    "source": [
     "# NAS (Neural Architecture Search) for PointPillars Lidar Detection on Vertex AI with TF-vision\n",
     "\n",
     "Make sure that you have read the [required documentations](https://cloud.google.com/vertex-ai/docs/training/neural-architecture-search/overview#reading_order)\n",
     "before executing this notebook.\n",
     "NOTE: This notebook is meant to run pre-built trainer code with pre-built search-spaces. If you want to run your own trainer\n",
     "code or create your own NAS search-space from scratch, then do not use this notebook.\n",
     "\n",
     "Expected best model result:\n",
     "- The mAP/mAPH on Waymo-Open-Dataset: 49.33/48.59\n",
     "- The inference latency on 1 V100 GPU: 57.3 ms\n",
     "\n",
     "The mAP/mAPH of our chosen [baseline model](https://github.com/tensorflow/models/tree/master/official/projects/pointpillars) is 45.96/45.35. The best model searched by NAS improved metrics by 10% without significant latency change. Note that the information is not supposed to used as a benchmark.\n",
     "This example is just meant to demonstrate NAS LiDAR search space against our chosen baseline.\n",
     "\n",
     "Experiment setup to reproduce the result:\n",
     "- Stage-1 search\n",
     "    - Number of trials: 1300\n",
     "    - Number of GPUs per trial: 4\n",
     "    - GPU type: TESLA_V100\n",
     "    - Average time per trial: 4 hours\n",
     "    - Number of parallel trials: 10\n",
     "    - GPU quota used: 40 V100 GPUs\n",
     "    - Time to run: 21.6 days\n",
     "    Since the number of days is more than 14 days, you will have to [resume the search job](https://cloud.google.com/vertex-ai/docs/training/neural-architecture-search/nas-client#nas-cli-search-resume).\n",
     "    If you have higher GPU quota, then the runtime will decrease proportionately.\n",
     "    - GPU hours: 20,800 V100 GPU hours\n",
     "- Stage-2 full-training with top 5 models\n",
     "    - Number of trials: 5\n",
     "    - Number of GPUs per trial: 8\n",
     "    - GPU type: TESLA_V100\n",
     "    - Average time per trial: 12 days\n",
     "    - Number of parallel trials: 5\n",
     "    - GPU quota used: 40 V100 GPUs\n",
     "    - Time to run: 12 days\n",
     "    - GPU hours: 11,520 V100 GPU hours\n",
     "\n",
     "Stage1 search cost: ~$85,000\n",
     "Stage2 full-training cost: ~$45,000\n",
     "Total cost: ~$130,000\n",
     "\n",
     "\n",
     "Here are the **pre-requisites** before you can start using this notebook: \n",
     "1. Your GCP project should have been (a) clear-listed and (b) a GPU quota should have been allocated for the NAS jobs, refert to [this](https://cloud.google.com/vertex-ai/docs/training/neural-architecture-search/environment-setup#device-quota) for requesting GPU quota. This notebook requires 200 GPUs to run experiments.\n",
     "2. You have selected a python3 kernel to run this notebook."
    ]
   },
   {
    "cell_type": "markdown",
    "metadata": {},
    "source": [
     "# Install required libraries"
    ]
   },
   {
    "cell_type": "code",
    "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
     "%%sh\n",
     "\n",
     "# Libraries required for using the NAS client.\n",
     "pip install pyglove==0.1.0\n",
     "\n",
     "# Libraries required for pre-processing Waymo-Open-Dataset. \n",
     "pip install keras==2.6.0\n",
     "pip install tensorflow==2.6.0\n",
     "pip install tf-models-official==2.7.2\n",
     "pip install waymo-open-dataset-tf-2-6-0\n",
     "pip install apache-beam[gcp]==2.42.0 --user"
    ]
   },
   {
    "cell_type": "markdown",
    "metadata": {},
    "source": [
     "**NOTE: Please restart the notebook after installing above libraries successfully.**"
    ]
   },
   {
    "cell_type": "markdown",
    "metadata": {},
    "source": [
     "# Download source code\n",
     "This needs to be done just once.\n"
    ]
   },
   {
    "cell_type": "code",
    "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
     "%%sh\n",
     "\n",
     "# NOTE: It is ok for this step to fail if the directory exists.\n",
     "mkdir ~/nas_experiment"
    ]
   },
   {
    "cell_type": "code",
    "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
     "%%sh\n",
     "rm -r -f ~/nas_experiment/nas_codes\n",
     "git clone https://github.com/google/vertex-ai-nas.git ~/nas_experiment/nas_codes"
    ]
   },
   {
    "cell_type": "markdown",
    "metadata": {},
    "source": [
     "# Set code path"
    ]
   },
   {
    "cell_type": "code",
    "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
     "import os\n",
     "os.chdir('/home/jupyter/nas_experiment/nas_codes')"
    ]
   },
   {
    "cell_type": "markdown",
    "metadata": {},
    "source": [
     "# Set up environment variables\n",
     "Here we set up the environment variables.\n",
     "NOTE: These have to be set-up every time you run a new session because the later code-blocks use them.\n"
    ]
   },
   {
    "cell_type": "code",
    "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
     "# Set a unique USER name. This is used for creating a unique job-name identifier.\n",
     "%env USER=<fill>\n",
     "# Set a region to launch jobs into.\n",
     "%env REGION=<fill>\n",
     "# Set any unique docker-id for this run. When the next section builds a docker, then this id will be used to tag it.\n",
     "%env TRAINER_DOCKER_ID=<fill>\n",
     "# The GCP project-id must be the one that has been clear-listed for the NAS jobs. \n",
     "%env PROJECT_ID=<fill>\n",
     "# Set an output working directory for the NAS jobs. The GCP project should have write access to \n",
     "# this GCS output directory. A simple way to ensure this is to use a bucket inside the same GCP project.\n",
     "# NOTE: The region of the bucket must be the same as job's.\n",
     "%env GCS_ROOT_DIR=<fill>\n",
     "# Set the accelerator device type.\n",
     "%env DEVICE=<fill>"
    ]
   },
   {
    "cell_type": "markdown",
    "metadata": {},
    "source": [
     "**NOTE:** The following set up steps need to be done just once."
    ]
   },
   {
    "cell_type": "code",
    "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
     "%%sh\n",
     "\n",
     "# Authenticate docker for your artifact registry.\n",
     "gcloud auth configure-docker ${REGION}-docker.pkg.dev"
    ]
   },
   {
    "cell_type": "code",
    "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
     "%%sh\n",
     "\n",
     "# NOTE: This needs to be just once for the first time. It is ok for this to FAIL if the GCS bucket already exists.\n",
     "\n",
     "# Create the output directory. \n",
     "gsutil mkdir $GCS_ROOT_DIR"
    ]
   },
   {
    "cell_type": "markdown",
    "metadata": {},
    "source": [
     "# Prepare dataset"
    ]
   },
   {
    "cell_type": "markdown",
    "metadata": {},
    "source": [
     "We use [Waymo-Open-Dataset](https://waymo.com/open/) as the example dataset for training and evaluation. We run a script to pre-process the Waymo-Open-Dataset, which converts the raw [Lidar frame data](https://github.com/waymo-research/waymo-open-dataset/blob/master/waymo_open_dataset/dataset.proto#L370) to a format that can be fed into the model."
    ]
   },
   {
    "cell_type": "code",
    "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
     "%%sh\n",
     "\n",
     "# Set the GCS path of source Waymo-Open-Dataset. The GCP project should have read access to the data-location.\n",
     "SRC_DIR='gs://waymo_open_dataset_v_1_2_0_individual_files'\n",
     "# Set the GCS path of processed dataset. The GCP project should have write access to the data-location.\n",
     "DST_DIR=<fill>\n",
     "\n",
     "# Set the runner for beam. See https://beam.apache.org/documentation/#runners for distributed runners.\n",
     "RUNNER=\"DirectRunner\"\n",
     "\n",
     "export PYTHONPATH=\"/home/jupyter/nas_experiment/nas_codes/:$PYTHONPATH\"\n",
     "python3 tf_vision/pointpillars/tools/process_wod.py \\\n",
     "--src_dir=${SRC_DIR} \\\n",
     "--dst_dir=${DST_DIR} \\\n",
     "--pipeline_options=\"--runner=${RUNNER}\""
    ]
   },
   {
    "cell_type": "code",
    "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
     "# Set the GCS paths to the training and validation datasets. The GCP project should have read access to the data-location.\n",
     "%env STAGE1_TRAINING_DATA_PATH=<fill>\n",
     "%env STAGE1_VALIDATION_DATA_PATH=<fill>\n",
     "%env STAGE2_TRAINING_DATA_PATH=<fill>\n",
     "%env STAGE2_VALIDATION_DATA_PATH=<fill>"
    ]
   },
   {
    "cell_type": "markdown",
    "metadata": {},
    "source": [
     "# Validate and Visualize data format"
    ]
   },
   {
    "cell_type": "markdown",
    "metadata": {},
    "source": [
     "The following code verifies that the data can be loaded properly before you run the experiments."
    ]
   },
   {
    "cell_type": "code",
    "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
     "%matplotlib inline\n",
     "import os\n",
     "import matplotlib as mpl\n",
     "import matplotlib.pyplot as plt\n",
     "import matplotlib.image as mm\n",
     "import matplotlib.patches as patches\n",
     "import numpy as np\n",
     "import tensorflow as tf\n",
     "from tf_vision.pointpillars.configs import pointpillars as cfg\n",
     "from tf_vision.pointpillars.dataloaders import decoders\n",
     "\n",
     "def draw_labels(plt, ax, example):\n",
     "  classes = example['gt_classes'].numpy()\n",
     "  boxes = example['gt_boxes'].numpy()\n",
     "  headings = example['gt_attributes']['heading'].numpy()\n",
     "  for i in range(len(classes)):\n",
     "    xmin, xmax, ymin, ymax, heading, clss = \\\n",
     "        boxes[i][1], boxes[i][3], boxes[i][0], boxes[i][2], headings[i][0], classes[i]\n",
     "    length = xmax - xmin\n",
     "    width = ymax - ymin\n",
     "    center_x = xmin + length * 0.5\n",
     "    center_y = ymin + width * 0.5\n",
     "\n",
     "    color = 'red'\n",
     "    if clss == 1: color = 'green'\n",
     "    elif clss == 2: color = 'red'\n",
     "    elif clss == 3: color = 'orange'\n",
     "    t = mpl.transforms.Affine2D().rotate_around(center_x, center_y, -heading) + ax.transData\n",
     "\n",
     "    rect = patches.Rectangle((xmin, ymin), length, width, linewidth=4, edgecolor=color, facecolor='none')\n",
     "    rect.set_transform(t)\n",
     "    ax.add_patch(rect)\n",
     "    plt.scatter(center_x, center_y, s=100, c=color, marker='o', linewidths=3)\n",
     "\n",
     "def draw_example(example, height, width):\n",
     "  pillars = example['pillars'].numpy()\n",
     "  indices = example['indices'].numpy()\n",
     "  fig, ax = plt.subplots(figsize=(height*0.08, width*0.08))\n",
     "  img = np.zeros([height, width])\n",
     "\n",
     "  for i in range(len(pillars)):\n",
     "    index = indices[i]\n",
     "    img[index[0]][index[1]] = 1\n",
     "  plt.imshow(img)\n",
     "\n",
     "  draw_labels(plt, ax, example)\n",
     "  plt.show()\n",
     "\n",
     "\n",
     "pillars_config = cfg.PillarsConfig()\n",
     "image_config = cfg.ImageConfig()\n",
     "decoder = decoders.ExampleDecoder(image_config, pillars_config)\n",
     "dataset_files = tf.io.gfile.glob(os.environ['STAGE1_TRAINING_DATA_PATH'])\n",
     "dataset = tf.data.TFRecordDataset(dataset_files[2], compression_type='GZIP')\n",
     "for example in dataset:\n",
     "  frame = decoder.decode(example)\n",
     "  draw_example(frame, image_config.height, image_config.width)\n",
     "  break"
    ]
   },
   {
    "cell_type": "markdown",
    "metadata": {},
    "source": [
     "# Build container\n",
     "The container must be built the first time and then every time the source-code is modified. Otherwise, there is no need to run this step. This step internally builds the *Dockerfile* in the source-code directory and then pushes the docker to the cloud. "
    ]
   },
   {
    "cell_type": "code",
    "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
     "%%sh\n",
     "\n",
     "# NOTE: This step can take several minutes when run for the first time.\n",
     "\n",
     "python3 vertex_nas_cli.py build \\\n",
     "--project_id=${PROJECT_ID} \\\n",
     "--trainer_docker_id=${TRAINER_DOCKER_ID} \\\n",
     "--region=${REGION} \\\n",
     "--trainer_docker_file=\"tf_vision/pointpillars.Dockerfile\""
    ]
   },
   {
    "cell_type": "markdown",
    "metadata": {},
    "source": [
     "# Launch NAS stage 1 job\n",
     "If you want to customize this notebook for your own dataset, then you must\n",
     "read the [required documentations](https://cloud.google.com/vertex-ai/docs/training/neural-architecture-search/overview#reading_order)\n",
     "to ensure that you set up the proxy-task and other settings properly."
    ]
   },
   {
    "cell_type": "code",
    "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
     "%%sh\n",
     "\n",
     "DATE=\"$(date '+%Y%m%d_%H%M%S')\"\n",
     "JOB_ID=\"${USER}_nas_pointpillars_${DATE}\"\n",
     "\n",
     "CMD=\"\n",
     "python3 vertex_nas_cli.py search \\\n",
     "--project_id=${PROJECT_ID} \\\n",
     "--region=${REGION} \\\n",
     "--trainer_docker_id=${TRAINER_DOCKER_ID} \\\n",
     "--job_name=${JOB_ID} \\\n",
     "--max_nas_trial=3000 \\\n",
     "--max_parallel_nas_trial=50 \\\n",
     "--max_failed_nas_trial=300 \\\n",
     "--use_prebuilt_trainer=True \\\n",
     "--prebuilt_search_space=\"pointpillars\" \\\n",
     "--accelerator_type=${DEVICE} \\\n",
     "--num_gpus=4 \\\n",
     "--master_machine_type=\"n1-highmem-32\" \\\n",
     "--root_output_dir=${GCS_ROOT_DIR} \\\n",
     "--search_docker_flags \\\n",
     "params_override=\"tf_vision/configs/experiments/pointpillars_search_gpu.yaml\" \\\n",
     "training_data_path=${STAGE1_TRAINING_DATA_PATH} \\\n",
     "validation_data_path=${STAGE1_VALIDATION_DATA_PATH} \\\n",
     "model=\"pointpillars\" \\\n",
     "multiple_eval_during_search=false\n",
     "\"\n",
     "\n",
     "echo Executing command: ${CMD}\n",
     "    \n",
     "${CMD}"
    ]
   },
   {
    "cell_type": "markdown",
    "metadata": {},
    "source": [
     "# Inspect NAS search progress\n",
     "A periodic evaluation while the search is going on can help decide if the search job has converged. This code-block shows how to generate summary of top N trials so far."
    ]
   },
   {
    "cell_type": "code",
    "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
     "# Set the stage1 search-job id. It's a numeric value returned by the Vertex service.\n",
     "%env JOB_ID=<fill>"
    ]
   },
   {
    "cell_type": "code",
    "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
     "%%sh\n",
     "\n",
     "mkdir -p /home/jupyter/nas_experiment/jobs\n",
     "\n",
     "OUTPUT_FILE=/home/jupyter/nas_experiment/jobs/${JOB_ID}.yaml\n",
     "\n",
     "python3 vertex_nas_cli.py list_trials \\\n",
     "--project_id=${PROJECT_ID} \\\n",
     "--job_id=${JOB_ID} \\\n",
     "--region=${REGION} \\\n",
     "--trials_output_file=${OUTPUT_FILE}\n",
     "\n",
     "cat ${OUTPUT_FILE}"
    ]
   },
   {
    "cell_type": "markdown",
    "metadata": {},
    "source": [
     "# Launch NAS stage 2 job"
    ]
   },
   {
    "cell_type": "code",
    "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
     "%%sh\n",
     "\n",
     "DATE=\"$(date '+%Y%m%d_%H%M%S')\"\n",
     "\n",
     "# TRIAL_ID is one of the best performing trials which has to be finetuned.\n",
     "TRIAL_IDS=<fill> # The top trials chosen for training to converge.\n",
     "\n",
     "CMD=\"\n",
     "\n",
     "python3 vertex_nas_cli.py train \\\n",
     "--project_id=${PROJECT_ID} \\\n",
     "--region=${REGION} \\\n",
     "--trainer_docker_id=${TRAINER_DOCKER_ID} \\\n",
     "--use_prebuilt_trainer=True \\\n",
     "--prebuilt_search_space=\"pointpillars\" \\\n",
     "--train_accelerator_type=${DEVICE} \\\n",
     "--train_num_gpus=8 \\\n",
     "--train_master_machine_type=\"n1-highmem-32\" \\\n",
     "--root_output_dir=${GCS_ROOT_DIR} \\\n",
     "--search_job_id=${JOB_ID} \\\n",
     "--search_job_region=${REGION} \\\n",
     "--train_nas_trial_numbers=${TRIAL_IDS} \\\n",
     "--train_job_suffix=\"stage2_${DATE}\" \\\n",
     "--train_docker_flags \\\n",
     "params_override=\"tf_vision/configs/experiments/pointpillars_search_finetune_gpu.yaml\" \\\n",
     "training_data_path=${STAGE2_TRAINING_DATA_PATH} \\\n",
     "validation_data_path=${STAGE2_VALIDATION_DATA_PATH} \\\n",
     "model=\"pointpillars\"\n",
     "\"\n",
     "\n",
     "echo Executing command: ${CMD}\n",
     "    \n",
     "${CMD}"
    ]
   }
  ],
  "metadata": {
   "environment": {
    "kernel": "python3",
    "name": "common-cu110.m100",
    "type": "gcloud",
    "uri": "gcr.io/deeplearning-platform-release/base-cu110:m100"
   },
   "kernelspec": {
    "display_name": "Python 3 (ipykernel)",
    "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": 4
 }