notebooks/vertex_nas_classification_tfvision.ipynb

{
  "cells": [
   {
    "cell_type": "markdown",
    "metadata": {},
    "source": [
     "# NAS (Neural Architecture Search) for Classification 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",
     "This notebook shows example of [MnasNet](https://arxiv.org/abs/1807.11626) paper result on Imagenet data.\n",
     "According to the paper, MnasNet achieves 75.2% top-1 accuracy with 78ms latency on a Pixel phone, \n",
     "which is 1.8x faster than MobileNetV2 with 0.5% higher accuracy and 2.3x faster than NASNet with 1.2% higher accuracy.\n",
     "However, this notebook uses GPUs instead of TPUs for training and uses cloud-CPU (n1-highmem-8) to evaluate latency.\n",
     "With this notebook, the expected Stage2 top-1 accuracy on MNasnet is 75.2% with 50ms latency on cloud-CPU (n1-highmem-8).\n",
     "The detailed settings for this notebook are:\n",
     "- Stage-1 search\n",
     "    - Number of trials: 2000\n",
     "    - Number of GPUs per trial: 2\n",
     "    - GPU type: TESLA_T4\n",
     "    - Avg single trial training time: 3 hours\n",
     "    - Number of parallel trials: 10\n",
     "    - GPU quota used: 10*2 = 20 T4 GPUs\n",
     "    - Time to run: 25 days\n",
     "    If you have higher GPU quota, then the runtime will decrease proportionately.\n",
     "    - GPU hours: 12000 T4 GPU hours\n",
     "- Stage-2 full-training with top 10 models\n",
     "    - Number of trials: 10\n",
     "    - Number of GPUs per trial: 4\n",
     "    - GPU type: TESLA_T4\n",
     "    - Avg single trial training time: ~9 days\n",
     "    - Number of parallel trials: 10\n",
     "    - GPU quota used: 10*4 = 40 T4 GPUs\n",
     "    You can also run this with just 20 T4 GPUs by running the job twice with just 5 models at\n",
     "    a time instead of all 10 in parallel.\n",
     "    - Time to run: ~9 days\n",
     "    - GPU hours: 8960 T4 GPU hours\n",
     "\n",
     "Stage1 search cost: ~$15,000\n",
     "Stage2 full-training cost: ~$8,000\n",
     "Total cost: ~$23,000\n",
     "\n",
     "\n",
     "You can also test drive MnasNet with just few trials with much lower cost.\n",
     "See [here for the test drive settings and cost](https://cloud.google.com/vertex-ai/docs/training/neural-architecture-search/overview#mnasnet_test_drive).\n",
     "For this use case, run this notebook only till section 'Test drive only: Launch NAS stage 1 job with latency constraint'.\n",
     "\n",
     "\n",
     "Here are the **pre-requisites** before you can start using this notebook: \n",
     "1. Your GCP project should have been (a) allow-listed and (b) [a GPU quota should have been allocated](https://cloud.google.com/vertex-ai/docs/training/neural-architecture-search/environment-setup#device-quota) for the NAS jobs.\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",
     "pip install tensorflow==2.7.0 --user\n",
     "pip install tf-models-official==2.7.1\n",
     "pip install pyglove==0.1.0"
    ]
   },
   {
    "cell_type": "markdown",
    "metadata": {},
    "source": [
     "**NOTE: Please restart the notebook after installing above libraries successfully.**"
    ]
   },
   {
    "cell_type": "markdown",
    "metadata": {},
    "source": [
     "# Download source code\n",
     "\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 -p ~/nas_experiment"
    ]
   },
   {
    "cell_type": "code",
    "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
     "%%sh\n",
     "\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(os.path.join('/home/jupyter/nas_experiment/nas_codes'))"
    ]
   },
   {
    "cell_type": "markdown",
    "metadata": {},
    "source": [
     "# Set up environment variables\n",
     "\n",
     "Here we set up the environment variables.\n",
     "\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",
     "# If you only want to test-drive and do not have enough GPU quota, then you can use 'us-central1' region\n",
     "# which should have a default quota of 12 Nvidia T4 GPUs.\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",
     "%env LATENCY_CALCULATOR_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=NVIDIA_TESLA_T4\n",
     "\n",
     "# Set the GCS paths to the training and validation datasets. The GCP project should have read access to the data-location.\n",
     "# Please read the \"Data-Preparation\" section \n",
     "# (https://cloud.google.com/vertex-ai/docs/training/neural-architecture-search/pre-built-trainer#data-preparation)\n",
     "# in the documentation to ensure that the data is in an appropriate format\n",
     "# suitable for the NAS pipeline. The documentation also mentions how you can download and prepare the ImageNet dataset.\n",
     "# You can run the \"Validate and Visualize data format\" section in this notebook \n",
     "# to verify that the data can be loaded properly.\n",
     "# Update the path to ImageNet data below.\n",
     "%env STAGE1_TRAINING_DATA_PATH=<gs://path-to>/imagenet/train-00[0-8]??-of-01024\n",
     "%env STAGE1_VALIDATION_DATA_PATH=<gs://path-to>/imagenet/train-009??-of-01024,gs://cloud-nas-public-eu/classification/imagenet/train-01???-of-01024\n",
     "%env STAGE2_TRAINING_DATA_PATH=<gs://path-to>/imagenet/train*\n",
     "%env STAGE2_VALIDATION_DATA_PATH=<gs://path-to>/imagenet/validation*"
    ]
   },
   {
    "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": [
     "# Validate and Visualize data format\n",
     "\n",
     "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": [
     "import tensorflow as tf\n",
     "import matplotlib.pyplot as plt\n",
     "from official.vision.beta.dataloaders import classification_input\n",
     "\n",
     "dataset = tf.data.Dataset.list_files(os.environ.get('STAGE1_TRAINING_DATA_PATH'), shuffle=False).apply(tf.data.TFRecordDataset)\n",
     "dataset = dataset.map(classification_input.Decoder().decode).batch(1)\n",
     "\n",
     "num_examples = 4\n",
     "_, ax = plt.subplots(num_examples, 1, figsize=(100, 64))\n",
     "for (i, example) in enumerate(dataset.take(num_examples)):\n",
     "    image = tf.io.decode_image(example['image/encoded'][0], channels=3)\n",
     "    image.set_shape([None, None, 3])\n",
     "    image = image.numpy()\n",
     "    \n",
     "    ax[i].imshow(image)\n",
     "    ax[i].grid(False)\n",
     "    ax[i].set_title('label: {}'.format(example['image/class/label'][0].numpy()))"
    ]
   },
   {
    "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/nas_multi_trial.Dockerfile\" \\\n",
     "--latency_calculator_docker_id=${LATENCY_CALCULATOR_DOCKER_ID} \\\n",
     "--latency_calculator_docker_file=\"tf_vision/latency_computation_using_saved_model.Dockerfile\""
    ]
   },
   {
    "cell_type": "markdown",
    "metadata": {},
    "source": [
     "# Test drive only: Launch NAS stage 1 job with latency constraint\n",
     "If you do not want to run a full MNasNet run and only want to test drive with [just few trials as described here,](https://cloud.google.com/vertex-ai/docs/training/neural-architecture-search/overview#mnasnet_test_drive)\n",
     "then only run the following command and skip the rest of the notebook."
    ]
   },
   {
    "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_tfvision_icn_latency_${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=25 \\\n",
     "--max_parallel_nas_trial=6 \\\n",
     "--max_failed_nas_trial=10 \\\n",
     "--use_prebuilt_trainer=True \\\n",
     "--prebuilt_search_space=\"mnasnet\" \\\n",
     "--accelerator_type=${DEVICE} \\\n",
     "--num_gpus=2 \\\n",
     "--root_output_dir=${GCS_ROOT_DIR} \\\n",
     "--latency_calculator_docker_id=${LATENCY_CALCULATOR_DOCKER_ID} \\\n",
     "--target_device_type=CPU \\\n",
     "--use_prebuilt_latency_calculator=True \\\n",
     "--search_docker_flags \\\n",
     "params_override=\"tf_vision/configs/experiments/mnasnet_search_gpu.yaml\" \\\n",
     "training_data_path=${STAGE1_TRAINING_DATA_PATH} \\\n",
     "validation_data_path=${STAGE1_VALIDATION_DATA_PATH} \\\n",
     "model=\"classification\" \\\n",
     "target_device_latency_ms=50\n",
     "\"\n",
     "\n",
     "echo Executing command: ${CMD}\n",
     "    \n",
     "${CMD}"
    ]
   },
   {
    "cell_type": "markdown",
    "metadata": {},
    "source": [
     "# Launch NAS stage 1 job with latency constraint\n",
     "If you want to customize this notebook for your own dataset other than ImageNet, 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_tfvision_icn_latency_${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=2000 \\\n",
     "--max_parallel_nas_trial=10 \\\n",
     "--max_failed_nas_trial=400 \\\n",
     "--use_prebuilt_trainer=True \\\n",
     "--prebuilt_search_space=\"mnasnet\" \\\n",
     "--accelerator_type=${DEVICE} \\\n",
     "--num_gpus=2 \\\n",
     "--root_output_dir=${GCS_ROOT_DIR} \\\n",
     "--latency_calculator_docker_id=${LATENCY_CALCULATOR_DOCKER_ID} \\\n",
     "--target_device_type=CPU \\\n",
     "--use_prebuilt_latency_calculator=True \\\n",
     "--search_docker_flags \\\n",
     "params_override=\"tf_vision/configs/experiments/mnasnet_search_gpu.yaml\" \\\n",
     "training_data_path=${STAGE1_TRAINING_DATA_PATH} \\\n",
     "validation_data_path=${STAGE1_VALIDATION_DATA_PATH} \\\n",
     "model=\"classification\" \\\n",
     "target_device_latency_ms=50\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",
     "python3 vertex_nas_cli.py list_trials \\\n",
     "--project_id=${PROJECT_ID} \\\n",
     "--job_id=${JOB_ID} \\\n",
     "--region=${REGION} \\\n",
     "--trials_output_file=/home/jupyter/nas_experiment/jobs/${JOB_ID}.yaml\n",
     "\n",
     "cat /home/jupyter/nas_experiment/jobs/${JOB_ID}.yaml"
    ]
   },
   {
    "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",
     "# Please modify the \"JOB_ID\", \"TRIAL_ID\", and the finetuning config file before running.\n",
     "# JOB_ID is numeric value you can find from the job UI in Pantheon.\n",
     "JOB_ID=<fill>\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=\"mnasnet\" \\\n",
     "--train_accelerator_type=${DEVICE} \\\n",
     "--train_num_gpus=4 \\\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/mnasnet_search_finetune_gpu.yaml\" \\\n",
     "training_data_path=${STAGE2_TRAINING_DATA_PATH} \\\n",
     "validation_data_path=${STAGE2_VALIDATION_DATA_PATH} \\\n",
     "model=\"classification\"\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",
    "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
 }