fine-tune.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "name": "fine-tune.ipynb",
      "provenance": [],
      "collapsed_sections": [],
      "machine_shape": "hm",
      "include_colab_link": true
    },
    "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.8.2"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/sayakpaul/MLPMixer-jax2tf/blob/main/fine-tune.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Yx9kQrATLdy5"
      },
      "source": [
        "## Imports"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "UqO1f2Z7QoOC"
      },
      "source": [
        "from tensorflow import keras\n",
        "import tensorflow as tf\n",
        "import tensorflow_hub as hub\n",
        "\n",
        "import tensorflow_datasets as tfds\n",
        "\n",
        "tfds.disable_progress_bar()\n",
        "\n",
        "import matplotlib.pyplot as plt\n",
        "import numpy as np"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "yS9v0obPLgVy"
      },
      "source": [
        "## Model building utility"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "yZ0gsA41RVVM"
      },
      "source": [
        "def get_model(\n",
        "    handle=\"https://tfhub.dev/sayakpaul/mixer_b16_i1k_fe/1\", \n",
        "    num_classes=5,\n",
        "):\n",
        "    hub_layer = hub.KerasLayer(handle, trainable=True)\n",
        "\n",
        "    model = keras.Sequential(\n",
        "        [\n",
        "            keras.layers.InputLayer((224, 224, 3)),\n",
        "            hub_layer,\n",
        "            keras.layers.Dense(num_classes, activation=\"softmax\"),\n",
        "        ]\n",
        "    )\n",
        "\n",
        "    return model"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "jk5VZIafSxGv"
      },
      "source": [
        "get_model().summary()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ELO3it7aLk6E"
      },
      "source": [
        "## Data input pipeline\n",
        "\n",
        "Code is reused from the [official repository](https://github.com/google-research/vision_transformer/blob/main/vit_jax/input_pipeline.py)."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "hE2MxLkbEfa7"
      },
      "source": [
        "BATCH_SIZE = 64\n",
        "AUTO = tf.data.AUTOTUNE"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Ouk89SHNS9us"
      },
      "source": [
        "def make_dataset(dataset: tf.data.Dataset, train: bool, image_size: int = 224):\n",
        "    def preprocess(image, label):\n",
        "        # For training, do a random crop and horizontal flip.\n",
        "        if train:\n",
        "            channels = image.shape[-1]\n",
        "            begin, size, _ = tf.image.sample_distorted_bounding_box(\n",
        "                tf.shape(image),\n",
        "                tf.zeros([0, 0, 4], tf.float32),\n",
        "                area_range=(0.05, 1.0),\n",
        "                min_object_covered=0,\n",
        "                use_image_if_no_bounding_boxes=True,\n",
        "            )\n",
        "            image = tf.slice(image, begin, size)\n",
        "\n",
        "            image.set_shape([None, None, channels])\n",
        "            image = tf.image.resize(image, [image_size, image_size])\n",
        "            if tf.random.uniform(shape=[]) > 0.5:\n",
        "                image = tf.image.flip_left_right(image)\n",
        "\n",
        "        else:\n",
        "            image = tf.image.resize(image, [image_size, image_size])\n",
        "\n",
        "        image = (image - 127.5) / 127.5\n",
        "        return image, label\n",
        "\n",
        "    if train:\n",
        "        dataset = dataset.shuffle(BATCH_SIZE * 10)\n",
        "\n",
        "    return dataset.map(preprocess, AUTO).batch(BATCH_SIZE).prefetch(AUTO)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Hl3i3onrLtO6"
      },
      "source": [
        "## `tf_flowers` dataset"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "lHJFUqFZE-0F"
      },
      "source": [
        "train_dataset, val_dataset = tfds.load(\n",
        "    \"tf_flowers\", split=[\"train[:90%]\", \"train[90%:]\"], as_supervised=True\n",
        ")\n",
        "\n",
        "num_train = tf.data.experimental.cardinality(train_dataset)\n",
        "num_val = tf.data.experimental.cardinality(val_dataset)\n",
        "print(f\"Number of training examples: {num_train}\")\n",
        "print(f\"Number of validation examples: {num_val}\")"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "rnjfXf8RLzCp"
      },
      "source": [
        "### Prepare dataset"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "DRd8kkcMFxSw"
      },
      "source": [
        "train_dataset = make_dataset(train_dataset, True)\n",
        "val_dataset = make_dataset(val_dataset, False)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "HbeqUHVdLz5J"
      },
      "source": [
        "### Visualize"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "O8Ui5B8hGNXQ"
      },
      "source": [
        "sample_images, _ = next(iter(train_dataset))\n",
        "\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() + 1) / 2)\n",
        "    plt.axis(\"off\")\n",
        "plt.show()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "h1XNooD2L2IY"
      },
      "source": [
        "## Learning rate scheduling \n",
        "\n",
        "For fine-tuning the authors follow a warm-up + [cosine | linear] schedule as per the [official notebook](https://colab.research.google.com/github/google-research/vision_transformer/blob/linen/vit_jax.ipynb). "
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "3cFQQaQoGjuF"
      },
      "source": [
        "# Reference:\n",
        "# https://www.kaggle.com/ashusma/training-rfcx-tensorflow-tpu-effnet-b2\n",
        "\n",
        "\n",
        "class WarmUpCosine(tf.keras.optimizers.schedules.LearningRateSchedule):\n",
        "    def __init__(\n",
        "        self, learning_rate_base, total_steps, warmup_learning_rate, warmup_steps\n",
        "    ):\n",
        "        super(WarmUpCosine, self).__init__()\n",
        "\n",
        "        self.learning_rate_base = learning_rate_base\n",
        "        self.total_steps = total_steps\n",
        "        self.warmup_learning_rate = warmup_learning_rate\n",
        "        self.warmup_steps = warmup_steps\n",
        "        self.pi = tf.constant(np.pi)\n",
        "\n",
        "    def __call__(self, step):\n",
        "        if self.total_steps < self.warmup_steps:\n",
        "            raise ValueError(\"Total_steps must be larger or equal to warmup_steps.\")\n",
        "        learning_rate = (\n",
        "            0.5\n",
        "            * self.learning_rate_base\n",
        "            * (\n",
        "                1\n",
        "                + tf.cos(\n",
        "                    self.pi\n",
        "                    * (tf.cast(step, tf.float32) - self.warmup_steps)\n",
        "                    / float(self.total_steps - self.warmup_steps)\n",
        "                )\n",
        "            )\n",
        "        )\n",
        "\n",
        "        if self.warmup_steps > 0:\n",
        "            if self.learning_rate_base < self.warmup_learning_rate:\n",
        "                raise ValueError(\n",
        "                    \"Learning_rate_base must be larger or equal to \"\n",
        "                    \"warmup_learning_rate.\"\n",
        "                )\n",
        "            slope = (\n",
        "                self.learning_rate_base - self.warmup_learning_rate\n",
        "            ) / self.warmup_steps\n",
        "            warmup_rate = slope * tf.cast(step, tf.float32) + self.warmup_learning_rate\n",
        "            learning_rate = tf.where(\n",
        "                step < self.warmup_steps, warmup_rate, learning_rate\n",
        "            )\n",
        "        return tf.where(\n",
        "            step > self.total_steps, 0.0, learning_rate, name=\"learning_rate\"\n",
        "        )"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "1-fPavRyMH1X"
      },
      "source": [
        "## Training hyperparameters\n",
        "\n",
        "These have been referred from the official notebooks ([1](https://colab.research.google.com/github/google-research/vision_transformer/blob/linen/vit_jax.ipynb) and [2](https://colab.research.google.com/github/google-research/vision_transformer/blob/master/vit_jax_augreg.ipynb)). \n",
        "\n",
        "Differences:\n",
        "\n",
        "* No gradient accumulation\n",
        "* Lower batch size for demoing on a single GPU (64 as opposed to 512)"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "bgsWyyhgHAaB"
      },
      "source": [
        "EPOCHS = 8\n",
        "TOTAL_STEPS = int((num_train / BATCH_SIZE) * EPOCHS)\n",
        "WARMUP_STEPS = 10\n",
        "INIT_LR = 0.03\n",
        "WAMRUP_LR = 0.006\n",
        "\n",
        "print(TOTAL_STEPS)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "LppzZVdbHwPE"
      },
      "source": [
        "scheduled_lrs = WarmUpCosine(\n",
        "    learning_rate_base=INIT_LR,\n",
        "    total_steps=TOTAL_STEPS,\n",
        "    warmup_learning_rate=WAMRUP_LR,\n",
        "    warmup_steps=WARMUP_STEPS,\n",
        ")\n",
        "\n",
        "lrs = [scheduled_lrs(step) for step in range(TOTAL_STEPS)]\n",
        "plt.plot(lrs)\n",
        "plt.xlabel(\"Step\", fontsize=14)\n",
        "plt.ylabel(\"LR\", fontsize=14)\n",
        "plt.show()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "i__WUIfcMpDk"
      },
      "source": [
        "### Optimizer and loss function"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "yblMGEOjIiyP"
      },
      "source": [
        "optimizer = keras.optimizers.SGD(scheduled_lrs, clipnorm=1.0)\n",
        "loss = keras.losses.SparseCategoricalCrossentropy()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "apdPFM_TMsJg"
      },
      "source": [
        "## Model training and validation"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "wd5XVBMZJlu_"
      },
      "source": [
        "model = get_model()\n",
        "model.compile(loss=loss, optimizer=optimizer, metrics=[\"accuracy\"])\n",
        "\n",
        "history = model.fit(train_dataset, validation_data=val_dataset, epochs=EPOCHS)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "InPOQqfJK3gR"
      },
      "source": [
        "plt.figure(figsize=(7, 7))\n",
        "history = history.history\n",
        "\n",
        "plt.plot(history[\"loss\"], label=\"train_loss\")\n",
        "plt.plot(history[\"val_loss\"], label=\"val_loss\")\n",
        "plt.plot(history[\"accuracy\"], label=\"train_accuracy\")\n",
        "plt.plot(history[\"val_accuracy\"], label=\"val_accuracy\")\n",
        "\n",
        "plt.legend()\n",
        "plt.show()"
      ],
      "execution_count": null,
      "outputs": []
    }
  ]
}