data2vec_vision_image_classification.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "data2vec-vision-image-classification.ipynb",
      "provenance": [],
      "collapsed_sections": [],
      "authorship_tag": "ABX9TyPEmUxVC6HNz8s8TjfpN6cJ",
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/sayakpaul/TF-2.0-Hacks/blob/master/data2vec_vision_image_classification.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "In this notebook, we'll fine-tune a model pre-trained using data2vec. We'll leverage the `transformers` library for loading the pre-trained model and we'll then append a custom classification head for fine-tuning. For data handling we'll use `tf.data`.\n",
        "\n",
        "Checkout the data2vec paper [here](https://arxiv.org/abs/2202.03555) and the model documentation page [here](https://huggingface.co/docs/transformers/main/en/model_doc/data2vec)."
      ],
      "metadata": {
        "id": "uP9zYDHBxN7V"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Setup"
      ],
      "metadata": {
        "id": "RJOiIwQix0y2"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# Otherwise TF's data2vec vision model won't be available.\n",
        "# Date: May 06, 2022\n",
        "!pip install git+https://github.com/huggingface/transformers -q"
      ],
      "metadata": {
        "id": "SDfJH5HmfpWF"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Imports"
      ],
      "metadata": {
        "id": "rNoyeQxdx2SK"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import tensorflow as tf\n",
        "from tensorflow import keras \n",
        "from tensorflow.keras import layers\n",
        "\n",
        "from transformers import TFAutoModelForImageClassification, TFData2VecVisionModel\n",
        "from transformers import create_optimizer\n",
        "\n",
        "import tensorflow_datasets as tfds\n",
        "\n",
        "tfds.disable_progress_bar()\n",
        "\n",
        "import pandas as pd\n",
        "import numpy as np\n",
        "import matplotlib.pyplot as plt"
      ],
      "metadata": {
        "id": "3BGFcoWPf3ht"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Primary constants"
      ],
      "metadata": {
        "id": "cM4gh8Ubx3nV"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# Model\n",
        "IMAGE_SIZE = [224, 224]\n",
        "MODEL_ID = \"facebook/data2vec-vision-base\"  # pre-trained model from which to fine-tune\n",
        "\n",
        "# Data\n",
        "BATCH_SIZE = 16  # on Colab/GPU, a higher batch size may throw(OOM)\n",
        "\n",
        "# Dataset\n",
        "CLASSES = [\n",
        "    \"dandelion\",\n",
        "    \"daisy\",\n",
        "    \"tulips\",\n",
        "    \"sunflowers\",\n",
        "    \"roses\",\n",
        "]  # don't change the order\n",
        "\n",
        "# Other constants\n",
        "MEAN = tf.constant([0.485 * 255, 0.456 * 255, 0.406 * 255])  # imagenet mean\n",
        "STD = tf.constant([0.229 * 255, 0.224 * 255, 0.225 * 255])  # imagenet std\n",
        "AUTO = tf.data.AUTOTUNE"
      ],
      "metadata": {
        "id": "cJ9uU4iMgJFF"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Data related utilities"
      ],
      "metadata": {
        "id": "nXHBYes2x5Nf"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# Closely aligns with https://github.com/huggingface/transformers/blob/main/src/transformers/models/beit/feature_extraction_beit.py\n",
        "# in regards to resizing and normalization.\n",
        "\n",
        "\n",
        "def make_dataset(dataset: tf.data.Dataset, train: bool, image_size: int = IMAGE_SIZE):\n",
        "    def preprocess(image, label):\n",
        "        # for training, do augmentation\n",
        "        if train:\n",
        "            if tf.random.uniform(shape=[]) > 0.5:\n",
        "                image = tf.image.flip_left_right(image)\n",
        "        image = tf.image.resize(image, size=image_size, method=\"bicubic\")\n",
        "        image = (image - MEAN) / STD  # normalization\n",
        "        return image, label\n",
        "\n",
        "    if train:\n",
        "        dataset = dataset.shuffle(BATCH_SIZE * 10)\n",
        "\n",
        "    dataset = dataset.map(preprocess, AUTO).batch(BATCH_SIZE)\n",
        "    # Transpose because the `transformers` model has a leading channel dimension.\n",
        "    dataset = dataset.map(lambda x, y: (tf.transpose(x, [0, 3, 1, 2]), y), AUTO)\n",
        "    return dataset.prefetch(AUTO)"
      ],
      "metadata": {
        "id": "Vz_kqZWrgykI"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Load and visualize the `tf_flowers` dataset"
      ],
      "metadata": {
        "id": "ifd4d8GPx_nu"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "train_dataset, val_dataset = tfds.load(\n",
        "    \"tf_flowers\",\n",
        "    split=[\"train[:90%]\", \"train[90%:]\"],\n",
        "    as_supervised=True,\n",
        "    try_gcs=False,  # gcs_path is necessary for tpu,\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}\")"
      ],
      "metadata": {
        "id": "7zVtKpWIokJd"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "train_dataset = make_dataset(train_dataset, True)\n",
        "val_dataset = make_dataset(val_dataset, False)"
      ],
      "metadata": {
        "id": "M-KhjVX9okxd"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "sample_images, sample_labels = next(iter(train_dataset))\n",
        "\n",
        "plt.figure(figsize=(5 * 3, 3 * 3))\n",
        "for n in range(15):\n",
        "    ax = plt.subplot(3, 5, n + 1)\n",
        "    image = (sample_images[n].numpy().transpose(1, 2, 0) * STD + MEAN).numpy()\n",
        "    image = (image - image.min()) / (\n",
        "        image.max() - image.min()\n",
        "    )  # convert to [0, 1] for avoiding matplotlib warning\n",
        "    plt.imshow(image)\n",
        "    plt.title(CLASSES[sample_labels[n]])\n",
        "    plt.axis(\"off\")\n",
        "plt.tight_layout()\n",
        "plt.show()"
      ],
      "metadata": {
        "id": "TS9Taf-UonJc"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Model building utilities"
      ],
      "metadata": {
        "id": "vjV7yqO1yCz_"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "def get_model():\n",
        "    base_model = TFData2VecVisionModel.from_pretrained(MODEL_ID)\n",
        "    base_model.trainable = True\n",
        "\n",
        "    inputs = layers.Input((3, IMAGE_SIZE[0], IMAGE_SIZE[1]))\n",
        "    x = base_model(inputs, training=False)\n",
        "    outputs = layers.Dense(5)(x.last_hidden_state[:, 0])\n",
        "    model = keras.Model(inputs, outputs)\n",
        "\n",
        "    return model"
      ],
      "metadata": {
        "id": "Sw4h4jFVg-Up"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "model = get_model()\n",
        "model.summary()"
      ],
      "metadata": {
        "id": "s1mU4RHgmgN8"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Compile model and train it"
      ],
      "metadata": {
        "id": "fx7-5d-SyUQ8"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "num_train_epochs = 3\n",
        "learning_rate = 3e-5\n",
        "weight_decay_rate = 0.01\n",
        "num_warmup_steps = 0\n",
        "\n",
        "\n",
        "num_train_steps = (num_train // BATCH_SIZE) * num_train_epochs\n",
        "optimizer, lr_schedule = create_optimizer(\n",
        "    init_lr=learning_rate,\n",
        "    num_train_steps=num_train_steps,\n",
        "    weight_decay_rate=weight_decay_rate,\n",
        "    num_warmup_steps=num_warmup_steps,\n",
        ")"
      ],
      "metadata": {
        "id": "J0c-VsDEoXN1"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "model.compile(\n",
        "    optimizer=optimizer,\n",
        "    loss=keras.losses.SparseCategoricalCrossentropy(from_logits=True),\n",
        "    metrics=[\"accuracy\"],\n",
        ")"
      ],
      "metadata": {
        "id": "rIenATY_o7-Y"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "history = model.fit(train_dataset, validation_data=val_dataset, epochs=num_train_epochs)"
      ],
      "metadata": {
        "id": "fYZd0jEio_X-"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Visualize training progress"
      ],
      "metadata": {
        "id": "KwWFR6vYyavE"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "result = pd.DataFrame(history.history)\n",
        "fig, ax = plt.subplots(2, 1, figsize=(10, 10))\n",
        "result[[\"accuracy\", \"val_accuracy\"]].plot(xlabel=\"epoch\", ylabel=\"score\", ax=ax[0])\n",
        "result[[\"loss\", \"val_loss\"]].plot(xlabel=\"epoch\", ylabel=\"score\", ax=ax[1])"
      ],
      "metadata": {
        "id": "R10teItxqSFs"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## References\n",
        "\n",
        "* https://github.com/sayakpaul/cait-tf/blob/main/notebooks/finetune.ipynb\n",
        "* https://colab.research.google.com/drive/1M4CEl6Jgf4KgqLHb4_bi3acJvNNLZqZd"
      ],
      "metadata": {
        "id": "b88z7qyQtYvH"
      }
    }
  ]
}