train_bit_keras_tuner.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Find the Kaggle Kernel of this notebook [here](https://www.kaggle.com/spsayakpaul/train-bit-keras-tuner).\n",
    "\n",
    "This notebook runs hyperparameter-tuning on a teacher model (based on [BiT ResNet101x3](https://arxiv.org/abs/1912.11370)) to further train a student using function matching (proposed in [Knowledge distillation: A good teacher is patient and consistent](https://arxiv.org/abs/2106.05237)). You can find the distillation notebook [here](https://www.kaggle.com/spsayakpaul/funmatch-distillation). To run this notebook you would need to have a billing enabled GCP account to use a GCS Bucket."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Setup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "_cell_guid": "b1076dfc-b9ad-4769-8c92-a6c4dae69d19",
    "_uuid": "8f2839f25d086af736a60e9eeb907d3b93b6e0e5",
    "execution": {
     "iopub.execute_input": "2021-07-26T11:52:52.497493Z",
     "iopub.status.busy": "2021-07-26T11:52:52.496722Z",
     "iopub.status.idle": "2021-07-26T11:53:01.038258Z",
     "shell.execute_reply": "2021-07-26T11:53:01.037073Z",
     "shell.execute_reply.started": "2021-07-26T11:52:52.497360Z"
    }
   },
   "outputs": [],
   "source": [
    "from kaggle_secrets import UserSecretsClient\n",
    "user_secrets = UserSecretsClient()\n",
    "user_credential = user_secrets.get_gcloud_credential()\n",
    "user_secrets.set_tensorflow_credential(user_credential)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-07-26T11:53:10.231074Z",
     "iopub.status.busy": "2021-07-26T11:53:10.230702Z",
     "iopub.status.idle": "2021-07-26T11:53:19.138938Z",
     "shell.execute_reply": "2021-07-26T11:53:19.137419Z",
     "shell.execute_reply.started": "2021-07-26T11:53:10.231043Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[33mWARNING: Running pip as root will break packages and permissions. You should install packages reliably by using venv: https://pip.pypa.io/warnings/venv\u001b[0m\n"
     ]
    }
   ],
   "source": [
    "!pip install keras-tuner -q --user"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-07-26T11:53:23.116016Z",
     "iopub.status.busy": "2021-07-26T11:53:23.115614Z",
     "iopub.status.idle": "2021-07-26T11:53:24.044046Z",
     "shell.execute_reply": "2021-07-26T11:53:24.042812Z",
     "shell.execute_reply.started": "2021-07-26T11:53:23.115978Z"
    }
   },
   "outputs": [],
   "source": [
    "import tensorflow_hub as hub\n",
    "import keras_tuner as kt\n",
    "import tensorflow as tf\n",
    "\n",
    "from pprint import pprint\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "import re\n",
    "\n",
    "# This needs to be done in order for the Hub module to communicate.\n",
    "import os\n",
    "os.environ[\"TFHUB_CACHE_DIR\"] = \"gs://funmatch-tf/model-cache-dir\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "`gs://funmatch-tf` is the GCS Bucket I created beforehand. To proceed, you'd need to create a GCS Bucket with a universally unique name and replace `funmatch-tf` with it. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-07-26T11:53:27.901118Z",
     "iopub.status.busy": "2021-07-26T11:53:27.900735Z",
     "iopub.status.idle": "2021-07-26T11:53:33.817355Z",
     "shell.execute_reply": "2021-07-26T11:53:33.816537Z",
     "shell.execute_reply.started": "2021-07-26T11:53:27.901084Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Number of accelerators:  8\n"
     ]
    }
   ],
   "source": [
    "try: # Cetect TPUs\n",
    "    tpu = None\n",
    "    tpu = tf.distribute.cluster_resolver.TPUClusterResolver() # TPU detection\n",
    "    tf.config.experimental_connect_to_cluster(tpu)\n",
    "    tf.tpu.experimental.initialize_tpu_system(tpu)\n",
    "    strategy = tf.distribute.TPUStrategy(tpu)\n",
    "except ValueError: # Detect GPUs\n",
    "    strategy = tf.distribute.MirroredStrategy() \n",
    "\n",
    "print(\"Number of accelerators: \", strategy.num_replicas_in_sync)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Constants and hyperparameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-07-26T11:53:42.383130Z",
     "iopub.status.busy": "2021-07-26T11:53:42.382613Z",
     "iopub.status.idle": "2021-07-26T11:53:42.388209Z",
     "shell.execute_reply": "2021-07-26T11:53:42.387350Z",
     "shell.execute_reply.started": "2021-07-26T11:53:42.383098Z"
    }
   },
   "outputs": [],
   "source": [
    "BATCH_SIZE = 64 * strategy.num_replicas_in_sync\n",
    "BIGGER = 160\n",
    "RESIZE = 128\n",
    "CENTRAL_FRAC = 0.875\n",
    "AUTO = tf.data.AUTOTUNE\n",
    "\n",
    "SCHEDULE_LENGTH = 500\n",
    "SCHEDULE_BOUNDARIES = [200, 300, 400]\n",
    "SCHEDULE_LENGTH = (SCHEDULE_LENGTH * 512 / BATCH_SIZE) "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Data loading and input preprocessing\n",
    "\n",
    "To know how these TFRecords were created refer to [this notebook](https://colab.research.google.com/github/sayakpaul/FunMatch-Distillation/blob/main/tfrecords_pets37.ipynb). **Be sure to update the GCS paths.** "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-07-26T11:53:56.975175Z",
     "iopub.status.busy": "2021-07-26T11:53:56.974673Z",
     "iopub.status.idle": "2021-07-26T11:53:57.295024Z",
     "shell.execute_reply": "2021-07-26T11:53:57.293776Z",
     "shell.execute_reply.started": "2021-07-26T11:53:56.975142Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "['gs://funmatch-tf/train/train_pets37-0-128.tfrec',\n",
      " 'gs://funmatch-tf/train/train_pets37-1-128.tfrec',\n",
      " 'gs://funmatch-tf/train/train_pets37-10-128.tfrec',\n",
      " 'gs://funmatch-tf/train/train_pets37-11-128.tfrec',\n",
      " 'gs://funmatch-tf/train/train_pets37-12-128.tfrec']\n",
      "['gs://funmatch-tf/validation/validation_pets37-0-128.tfrec',\n",
      " 'gs://funmatch-tf/validation/validation_pets37-1-128.tfrec',\n",
      " 'gs://funmatch-tf/validation/validation_pets37-2-112.tfrec']\n",
      "['gs://funmatch-tf/test/test_pets37-0-128.tfrec',\n",
      " 'gs://funmatch-tf/test/test_pets37-1-128.tfrec',\n",
      " 'gs://funmatch-tf/test/test_pets37-10-128.tfrec',\n",
      " 'gs://funmatch-tf/test/test_pets37-11-128.tfrec',\n",
      " 'gs://funmatch-tf/test/test_pets37-12-128.tfrec']\n"
     ]
    }
   ],
   "source": [
    "# This comes from this repository https://github.com/GoogleCloudPlatform/training-data-analyst.\n",
    "def count_data_items(filenames):\n",
    "    n = [int(re.compile(r\"-([0-9]*)\\.\").search(filename).group(1)) for filename in filenames]\n",
    "    return np.sum(n)\n",
    "\n",
    "train_pattern = \"gs://funmatch-tf/train/*.tfrec\"\n",
    "train_filenames = tf.io.gfile.glob(train_pattern)\n",
    "val_pattern = \"gs://funmatch-tf/validation/*.tfrec\"\n",
    "val_filenames = tf.io.gfile.glob(val_pattern)\n",
    "test_pattern = \"gs://funmatch-tf/test/*.tfrec\"\n",
    "test_filenames = tf.io.gfile.glob(test_pattern)\n",
    "\n",
    "DATASET_NUM_TRAIN_EXAMPLES = count_data_items(train_filenames)\n",
    "STEPS_PER_EPOCH = 10\n",
    "\n",
    "pprint(train_filenames[:5])\n",
    "pprint(val_filenames[:5])\n",
    "pprint(test_filenames[:5])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-07-26T11:54:09.280683Z",
     "iopub.status.busy": "2021-07-26T11:54:09.279940Z",
     "iopub.status.idle": "2021-07-26T11:54:09.296078Z",
     "shell.execute_reply": "2021-07-26T11:54:09.294955Z",
     "shell.execute_reply.started": "2021-07-26T11:54:09.280613Z"
    }
   },
   "outputs": [],
   "source": [
    "# Function to read the TFRecords, segregate the images and labels.\n",
    "def read_tfrecord(example, train):\n",
    "    features = {\n",
    "        \"image\": tf.io.FixedLenFeature([], tf.string), \n",
    "        \"class\": tf.io.FixedLenFeature([], tf.int64)\n",
    "    }\n",
    "    \n",
    "    example = tf.io.parse_single_example(example, features)\n",
    "    image = tf.image.decode_jpeg(example[\"image\"], channels=3)\n",
    "    \n",
    "    if train:\n",
    "        image = augment(image)\n",
    "    else:\n",
    "        image = tf.image.central_crop(image, central_fraction=CENTRAL_FRAC)\n",
    "        image = tf.image.resize(image, (RESIZE, RESIZE))\n",
    "        \n",
    "    image = tf.reshape(image, (RESIZE, RESIZE, 3))\n",
    "    image = tf.cast(image, tf.float32) / 255.0  \n",
    "    class_label = tf.cast(example[\"class\"], tf.int32)\n",
    "    return (image, class_label)\n",
    "\n",
    "# Load the TFRecords and create tf.data.Dataset\n",
    "def load_dataset(filenames, train):\n",
    "    opt = tf.data.Options()\n",
    "    opt.experimental_deterministic = False\n",
    "    \n",
    "    if not train:\n",
    "        opt.experimental_distribute.auto_shard_policy = tf.data.experimental.AutoShardPolicy.DATA\n",
    "    \n",
    "    dataset = tf.data.TFRecordDataset(filenames, num_parallel_reads=AUTO) \n",
    "    dataset = dataset.map(lambda x: (read_tfrecord(x, train)), num_parallel_calls=AUTO)\n",
    "    dataset = dataset.with_options(opt)\n",
    "    return dataset\n",
    "\n",
    "# Augmentation motivated from here:\n",
    "# https://github.com/google-research/big_transfer/blob/master/colabs/big_transfer_tf2.ipynb.\n",
    "def augment(image):\n",
    "    image = tf.image.resize(image, (BIGGER, BIGGER))\n",
    "    image = tf.image.random_flip_left_right(image)\n",
    "    image = tf.image.random_crop(image, [RESIZE, RESIZE, 3])\n",
    "    return image\n",
    "\n",
    "# Batch, shuffle, and repeat the dataset and prefetch it\n",
    "# well before the current epoch ends\n",
    "def batch_dataset(filenames, train, batch_size=BATCH_SIZE):\n",
    "    dataset = load_dataset(filenames, train)\n",
    "    if train:\n",
    "        dataset = dataset.repeat(int(SCHEDULE_LENGTH * BATCH_SIZE / DATASET_NUM_TRAIN_EXAMPLES * STEPS_PER_EPOCH) + 1 + 50)\n",
    "        dataset = dataset.shuffle(BATCH_SIZE*10)\n",
    "    dataset = dataset.batch(batch_size)\n",
    "    dataset = dataset.prefetch(AUTO) \n",
    "    return dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-07-26T11:54:12.310209Z",
     "iopub.status.busy": "2021-07-26T11:54:12.309620Z",
     "iopub.status.idle": "2021-07-26T11:54:12.730437Z",
     "shell.execute_reply": "2021-07-26T11:54:12.729059Z",
     "shell.execute_reply.started": "2021-07-26T11:54:12.310161Z"
    }
   },
   "outputs": [],
   "source": [
    "training_dataset = batch_dataset(train_filenames, True)\n",
    "validation_dataset = batch_dataset(val_filenames, False)\n",
    "test_dataset = batch_dataset(test_filenames, False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-07-26T11:54:21.934535Z",
     "iopub.status.busy": "2021-07-26T11:54:21.934053Z",
     "iopub.status.idle": "2021-07-26T11:54:21.938507Z",
     "shell.execute_reply": "2021-07-26T11:54:21.937569Z",
     "shell.execute_reply.started": "2021-07-26T11:54:21.934452Z"
    }
   },
   "outputs": [],
   "source": [
    "# sample_images, _ = next(iter(training_dataset))\n",
    "# plt.figure(figsize=(10, 10))\n",
    "# for n in range(25):\n",
    "#     ax = plt.subplot(5, 5, n + 1)\n",
    "#     plt.imshow(sample_images[n].numpy())\n",
    "#     plt.axis(\"off\")\n",
    "# plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Model related utilities"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Set up a custom model class"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-07-26T11:54:40.694230Z",
     "iopub.status.busy": "2021-07-26T11:54:40.693874Z",
     "iopub.status.idle": "2021-07-26T11:54:40.701549Z",
     "shell.execute_reply": "2021-07-26T11:54:40.700120Z",
     "shell.execute_reply.started": "2021-07-26T11:54:40.694200Z"
    }
   },
   "outputs": [],
   "source": [
    "# Referenced from: https://github.com/google-research/big_transfer/blob/master/colabs/big_transfer_tf2.ipynb. \n",
    "class MyBiTModel(tf.keras.Model):\n",
    "    def __init__(self, num_classes, module):\n",
    "        super().__init__()\n",
    "\n",
    "        self.num_classes = num_classes\n",
    "        self.head = tf.keras.layers.Dense(num_classes, kernel_initializer=\"zeros\")\n",
    "        self.bit_model = module\n",
    "  \n",
    "    def call(self, images):\n",
    "        bit_embedding = self.bit_model(images)\n",
    "        return self.head(bit_embedding)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Utility function to sample a learning rate"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-07-26T12:03:40.077264Z",
     "iopub.status.busy": "2021-07-26T12:03:40.076862Z",
     "iopub.status.idle": "2021-07-26T12:03:40.083025Z",
     "shell.execute_reply": "2021-07-26T12:03:40.081663Z",
     "shell.execute_reply.started": "2021-07-26T12:03:40.077231Z"
    }
   },
   "outputs": [],
   "source": [
    "def get_lr(hp):\n",
    "    initial_lr = hp.Choice(\"learning_rate\", values=[0.003, 1e-4, 1e-5, 5e-5])\n",
    "    lr = (initial_lr * BATCH_SIZE / 512) * strategy.num_replicas_in_sync \n",
    "    return lr"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Model building and compiling utility"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-07-26T12:03:42.945670Z",
     "iopub.status.busy": "2021-07-26T12:03:42.945323Z",
     "iopub.status.idle": "2021-07-26T12:03:42.953967Z",
     "shell.execute_reply": "2021-07-26T12:03:42.953061Z",
     "shell.execute_reply.started": "2021-07-26T12:03:42.945641Z"
    }
   },
   "outputs": [],
   "source": [
    "def build_model(hp):\n",
    "    model_url = \"https://tfhub.dev/google/bit/m-r101x3/1\"\n",
    "    module = hub.KerasLayer(model_url, trainable=True)\n",
    "    model = MyBiTModel(num_classes=37, module=module)\n",
    "    \n",
    "    lr = get_lr(hp)\n",
    "    lr_schedule = tf.keras.optimizers.schedules.PiecewiseConstantDecay(boundaries=SCHEDULE_BOUNDARIES, \n",
    "                                                                       values=[lr, lr*0.1, lr*0.001, lr*0.0001])\n",
    "    optimizer = tf.keras.optimizers.SGD(learning_rate=lr_schedule, momentum=0.9)\n",
    "    loss_fn = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True)\n",
    "    \n",
    "    model.compile(optimizer=optimizer,\n",
    "              loss=loss_fn,\n",
    "              metrics=[\"accuracy\"])\n",
    "    return model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Running the search\n",
    "\n",
    "Take note of the GCS paths!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-07-26T12:18:31.634169Z",
     "iopub.status.busy": "2021-07-26T12:18:31.633664Z",
     "iopub.status.idle": "2021-07-26T13:05:18.568332Z",
     "shell.execute_reply": "2021-07-26T13:05:18.567374Z",
     "shell.execute_reply.started": "2021-07-26T12:18:31.634135Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Trial 3 Complete [00h 11m 49s]\n",
      "val_accuracy: 0.9361413419246674\n",
      "\n",
      "Best val_accuracy So Far: 0.94701087474823\n",
      "Total elapsed time: 00h 45m 44s\n"
     ]
    }
   ],
   "source": [
    "tuner = kt.RandomSearch(\n",
    "    build_model,\n",
    "    objective=\"val_accuracy\",\n",
    "    executions_per_trial=2,\n",
    "    max_trials=3,\n",
    "    overwrite=True,\n",
    "    directory=\"gs://funmatch-tf/keras_tuner\",\n",
    "    project_name=\"funmatch\",\n",
    "    distribution_strategy=strategy,\n",
    ")\n",
    "\n",
    "tuner.search(training_dataset, \n",
    "    steps_per_epoch=STEPS_PER_EPOCH, \n",
    "    epochs=45, \n",
    "    validation_data=validation_dataset,\n",
    "    callbacks=[tf.keras.callbacks.EarlyStopping(\"val_accuracy\", patience=3)])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-07-26T13:06:21.088117Z",
     "iopub.status.busy": "2021-07-26T13:06:21.087394Z",
     "iopub.status.idle": "2021-07-26T13:06:21.238448Z",
     "shell.execute_reply": "2021-07-26T13:06:21.237344Z",
     "shell.execute_reply.started": "2021-07-26T13:06:21.088057Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Results summary\n",
      "Results in gs://funmatch-tf/keras_tuner/funmatch\n",
      "Showing 10 best trials\n",
      "Objective(name='val_accuracy', direction='max')\n",
      "Trial summary\n",
      "Hyperparameters:\n",
      "learning_rate: 1e-05\n",
      "Score: 0.94701087474823\n",
      "Trial summary\n",
      "Hyperparameters:\n",
      "learning_rate: 0.0001\n",
      "Score: 0.9361413419246674\n",
      "Trial summary\n",
      "Hyperparameters:\n",
      "learning_rate: 0.003\n",
      "Score: 0.8192934989929199\n"
     ]
    }
   ],
   "source": [
    "tuner.results_summary()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-07-26T13:11:29.178507Z",
     "iopub.status.busy": "2021-07-26T13:11:29.178112Z",
     "iopub.status.idle": "2021-07-26T13:13:08.378943Z",
     "shell.execute_reply": "2021-07-26T13:13:08.377461Z",
     "shell.execute_reply.started": "2021-07-26T13:11:29.178473Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "8/8 [==============================] - 32s 2s/step - loss: 0.3495 - accuracy: 0.9000\n",
      "Test top-1 accuracy: 90.32%\n"
     ]
    }
   ],
   "source": [
    "best_model = tuner.get_best_models(num_models=1)\n",
    "_, accuracy = best_model[0].evaluate(test_dataset)\n",
    "print(f\"Test top-1 accuracy: {round(accuracy * 100, 2)}%\")"
   ]
  }
 ],
 "metadata": {
  "environment": {
   "name": "tf2-cpu.2-4.m69",
   "type": "gcloud",
   "uri": "gcr.io/deeplearning-platform-release/tf2-cpu.2-4:m69"
  },
  "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.8.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}