notebooks/vertex_pytorch_mnasnet_tutorial.ipynb

{
  "cells": [
   {
    "cell_type": "markdown",
    "metadata": {},
    "source": [
     "# NAS (Neural Architecture Search) Pytorch Based Example on Vertex AI Platform  \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 is only meant to demonstrate a PyTorch trainer example using multiple GPUs with GCS data on Vertex AI platform.\n",
     "It is not meant to be used as a benchmark for [MNasnet paper](https://arxiv.org/abs/1807.11626) performance. Refer to this [blog](https://cloud.google.com/blog/products/ai-machine-learning/efficient-pytorch-training-with-vertex-ai) for more details.\n",
     "It runs only stage-1 search with just 2 trials and without latency constraints. It does not perform stage2 evaluation.\n",
     "Please see [Tensorflow MNasnet notebook](https://github.com/google/vertex-ai-nas/blob/main/notebooks/vertex_nas_classification_tfvision.ipynb) for full workflow example.\n",
     "\n",
     "Experiment setup:\n",
     "- Stage-1 search\n",
     "    - Number of trials: 2\n",
     "    - Number of GPUs per trial: 2\n",
     "    - GPU type: TESLA_V100\n",
     "    - Avg single trial training time: 3 hours\n",
     "    - Number of parallel trials: 2\n",
     "    - GPU quota used: 2*2 = 4 V100 GPUs\n",
     "    - Time to run: 0.125 days\n",
     "    - GPU hours: 12 V100 GPU hours\n",
     "Stage1 search cost: ~$53\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 16 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",
     "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 the code from Cloud Repository\n",
     "**NOTE:** The following set up steps need to be done just once.\n",
     "\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",
     "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 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",
     "%env GCS_ROOT_DIR=<gs://fill>\n",
     "# Set the region to be same as for your bucket. For example, `us-central1`.\n",
     "%env REGION=<fill>\n",
     "# Set the accelerator device type, for example NVIDIA_TESLA_V100.\n",
     "%env DEVICE=NVIDIA_TESLA_V100"
    ]
   },
   {
    "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",
     "# Enable the container registry API for your project.\n",
     "gcloud services enable containerregistry.googleapis.com --project=${PROJECT_ID}"
    ]
   },
   {
    "cell_type": "code",
    "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
     "%%sh\n",
     "\n",
     "# NOTE: This needs to be just once for the first time.\n",
     "\n",
     "# Create the output directory. \n",
     "# NOTE: It is ok for this step to fail if the directory already exists.\n",
     "gsutil mkdir $GCS_ROOT_DIR"
    ]
   },
   {
    "cell_type": "markdown",
    "metadata": {},
    "source": [
     "# Prepare dataset"
    ]
   },
   {
    "cell_type": "markdown",
    "metadata": {},
    "source": [
     "You need to download the [ImageNet dataset](https://image-net.org/download.php) first. Then run the following commands to shard the dataset."
    ]
   },
   {
    "cell_type": "code",
    "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
     "%%sh\n",
     "\n",
     "# Shard training dataset\n",
     "python3 pytorch/classification/shard_imagenet.py \\\n",
     "--image_list_file=<path-to-imagenet>/train_list.txt \\\n",
     "--output_pattern=<gcs-path>/imagenet/train-%06d.tar\n",
     "\n",
     "# Shard validation dataset\n",
     "python3 pytorch/classification/shard_imagenet.py \\\n",
     "--image_list_file=<path-to-imagenet>/validation_list.txt \\\n",
     "--output_pattern=<gcs-path>/imagenet/validation-%06d.tar"
    ]
   },
   {
    "cell_type": "code",
    "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
     "# Path to sharded training data\n",
     "%env TRAIN_DATA_PATH=<gcs-path>/imagenet/train-{000000..000466}.tar\n",
     "# Path to sharded validation data\n",
     "%env EVAL_DATA_PATH=<gcs-path>/imagenet/validation-{000000..000021}.tar"
    ]
   },
   {
    "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 may take few minutes the first time.\n",
     "python3 vertex_nas_cli.py build \\\n",
     "--project_id=${PROJECT_ID} \\\n",
     "--region=${REGION} \\\n",
     "--trainer_docker_id=${TRAINER_DOCKER_ID} \\\n",
     "--trainer_docker_file=pytorch/classification/classification.Dockerfile"
    ]
   },
   {
    "cell_type": "markdown",
    "metadata": {},
    "source": [
     "### Launch NAS - Stage 1 job"
    ]
   },
   {
    "cell_type": "code",
    "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
     "%%sh\n",
     "\n",
     "# For proxy task set-up.\n",
     "MAX_NAS_TRIAL=2\n",
     "MAX_PARALLEL_NAS_TRIAL=2\n",
     "MAX_FAILED_NAS_TRIAL=2\n",
     "\n",
     "job_name=\"${USER}_pytorch_example_$(date '+%Y%m%d_%H%M%S')\"\n",
     "\n",
     "# Run in Cloud\n",
     "CMD=\"\n",
     "\n",
     "python3 vertex_nas_cli.py search \\\n",
     "--project_id=${PROJECT_ID} \\\n",
     "--region=${REGION} \\\n",
     "--trainer_docker_id=${TRAINER_DOCKER_ID} \\\n",
     "--search_space_module=pytorch.classification.search_space.mnasnet_search_space \\\n",
     "--job_name=${job_name} \\\n",
     "--root_output_dir=${GCS_ROOT_DIR} \\\n",
     "--accelerator_type=${DEVICE} \\\n",
     "--num_gpus=2 \\\n",
     "--max_nas_trial=${MAX_NAS_TRIAL} \\\n",
     "--max_parallel_nas_trial=${MAX_PARALLEL_NAS_TRIAL} \\\n",
     "--max_failed_nas_trial=${MAX_FAILED_NAS_TRIAL} \\\n",
     "--nas_target_reward_metric=top_1_accuracy \\\n",
     "--search_docker_flags \\\n",
     "config_file=pytorch/classification/mnasnet_search_gpu.yaml \\\n",
     "train_data_path=${TRAIN_DATA_PATH} \\\n",
     "eval_data_path=${EVAL_DATA_PATH}\n",
     "\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
 }