TokenLearner-V1.1.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "BtuUbHj0jt71"
   },
   "source": [
    "# Imports and Setups\n",
    "\n",
    "This implementation currently does not run on CPUs. It seems like there is a problem with running grouped convs on CPUs in Keras in general ([reference](https://github.com/tensorflow/tensorflow/issues/34024)). Therefore, please run this on GPUs."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "PmrKKlKZkbnY"
   },
   "outputs": [],
   "source": [
    "!pip install -q -U tensorflow-addons"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "CsyTdBuTimEZ"
   },
   "outputs": [],
   "source": [
    "import tensorflow as tf\n",
    "\n",
    "from tensorflow import keras\n",
    "from tensorflow.keras import layers\n",
    "import tensorflow_addons as tfa\n",
    "\n",
    "from datetime import datetime\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "\n",
    "import math"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "b8EaplNdlrzI"
   },
   "source": [
    "# Hyperparameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "QWUI-ueYluAj"
   },
   "outputs": [],
   "source": [
    "# DATA\n",
    "BATCH_SIZE = 256\n",
    "AUTO = tf.data.AUTOTUNE\n",
    "INPUT_SHAPE = (32, 32, 3)\n",
    "NUM_CLASSES = 10\n",
    "\n",
    "# OPTIMIZER\n",
    "LEARNING_RATE = 1e-3\n",
    "WEIGHT_DECAY = 1e-4\n",
    "\n",
    "# TRAINING\n",
    "EPOCHS = 20\n",
    "\n",
    "# AUGMENTATION\n",
    "IMAGE_SIZE = 48  # We will resize input images to this size.\n",
    "PATCH_SIZE = 6  # Size of the patches to be extracted from the input images.\n",
    "NUM_PATCHES = (IMAGE_SIZE // PATCH_SIZE) ** 2\n",
    "\n",
    "# ViT ARCHITECTURE\n",
    "LAYER_NORM_EPS = 1e-6\n",
    "PROJECTION_DIM = 128\n",
    "NUM_HEADS = 4\n",
    "NUM_LAYERS = 4\n",
    "MLP_UNITS = [\n",
    "    PROJECTION_DIM * 2,\n",
    "    PROJECTION_DIM,\n",
    "]\n",
    "\n",
    "# TOKENLEARNER\n",
    "NUM_GROUPS = 4\n",
    "NUM_TOKENS = 8"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "o6iaIEIJkqP9"
   },
   "outputs": [],
   "source": [
    "# Get the CIFAR-10 dataset.\n",
    "(x_train, y_train), (x_test, y_test) = keras.datasets.cifar10.load_data()\n",
    "(x_train, y_train), (x_val, y_val) = (\n",
    "    (x_train[:40000], y_train[:40000]),\n",
    "    (x_train[40000:], y_train[40000:]),\n",
    ")\n",
    "print(f\"Training samples: {len(x_train)}\")\n",
    "print(f\"Validation samples: {len(x_val)}\")\n",
    "print(f\"Testing samples: {len(x_test)}\")\n",
    "\n",
    "train_ds = tf.data.Dataset.from_tensor_slices((x_train, y_train))\n",
    "train_ds = train_ds.shuffle(BATCH_SIZE * 100).batch(BATCH_SIZE).prefetch(AUTO)\n",
    "\n",
    "val_ds = tf.data.Dataset.from_tensor_slices((x_val, y_val))\n",
    "val_ds = val_ds.batch(BATCH_SIZE).prefetch(AUTO)\n",
    "\n",
    "test_ds = tf.data.Dataset.from_tensor_slices((x_test, y_test))\n",
    "test_ds = test_ds.batch(BATCH_SIZE).prefetch(AUTO)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "3gZwShY42gM2"
   },
   "source": [
    "# Augmentation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "uGTtuVR9nE3g"
   },
   "outputs": [],
   "source": [
    "data_augmentation = keras.Sequential(\n",
    "    [\n",
    "        layers.Rescaling(1 / 255.0),\n",
    "        layers.Resizing(INPUT_SHAPE[0] + 20, INPUT_SHAPE[0] + 20),\n",
    "        layers.RandomCrop(IMAGE_SIZE, IMAGE_SIZE),\n",
    "        layers.RandomFlip(\"horizontal\"),\n",
    "    ],\n",
    "    name=\"data_augmentation\",\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "2Ob3yvVA3E9u"
   },
   "source": [
    "# Patch Encoder"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "_XZH_aS4nqEE"
   },
   "outputs": [],
   "source": [
    "def position_embedding(\n",
    "    projected_patches, num_patches=NUM_PATCHES, projection_dim=PROJECTION_DIM\n",
    "):\n",
    "    # Build the positions.\n",
    "    positions = tf.range(start=0, limit=num_patches, delta=1)\n",
    "\n",
    "    # Encode the positions with an Embedding layer.\n",
    "    encoded_positions = layers.Embedding(\n",
    "        input_dim=num_patches, output_dim=projection_dim\n",
    "    )(positions)\n",
    "\n",
    "    # Add encoded positions to the projected patches.\n",
    "    return projected_patches + encoded_positions"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "TuUfoaGpoUxN"
   },
   "source": [
    "# MLP"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "mgr25Z2qoQb7"
   },
   "outputs": [],
   "source": [
    "def mlp(x, dropout_rate, hidden_units):\n",
    "    # Iterate over the hidden units and\n",
    "    # add Dense => Dropout.\n",
    "    for units in hidden_units:\n",
    "        x = layers.Dense(units, activation=tf.nn.gelu)(x)\n",
    "        x = layers.Dropout(dropout_rate)(x)\n",
    "    return x"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "PfXKWaNCogWJ"
   },
   "source": [
    "# TokenLearner (V1.1)\n",
    "\n",
    "From here on these symbols mean:\n",
    "- B: Batch Size\n",
    "- H: Height of the input\n",
    "- W: Width of the input\n",
    "- C: Channel size of the input\n",
    "\n",
    "As the authors note [here](https://github.com/google-research/scenic/blob/main/scenic/projects/token_learner/model.py#L104):\n",
    "\n",
    "> Instead of using 4 conv. layers with small channels to implement spatial\n",
    "  attention, this version uses 2 grouped conv. layers with more channels. It\n",
    "  also uses softmax instead of sigmoid. We confirmed that this version works\n",
    "  better when having limited training data, such as training with ImageNet1K\n",
    "  from scratch."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "Oz10FLn6AsGg"
   },
   "outputs": [],
   "source": [
    "def token_learner_v11(inputs, number_of_tokens=NUM_TOKENS):\n",
    "    # Layer normalize the inputs.\n",
    "    x = layers.LayerNormalization(epsilon=LAYER_NORM_EPS)(inputs)  # (B, H, W, C)\n",
    "\n",
    "    # Applying Conv2D => Reshape => Permute\n",
    "    # The reshape and permute is done to help with the next steps of\n",
    "    # multiplication and Global Average Pooling.\n",
    "    attention_maps = keras.Sequential(\n",
    "        [\n",
    "            layers.Conv2D(\n",
    "                filters=inputs.shape[-1],\n",
    "                kernel_size=(1, 1),\n",
    "                strides=(1, 1),\n",
    "                groups=NUM_GROUPS,\n",
    "                padding=\"same\",\n",
    "                use_bias=False,\n",
    "            ),\n",
    "            layers.Conv2D(\n",
    "                filters=number_of_tokens,\n",
    "                kernel_size=(1, 1),\n",
    "                strides=(1, 1),\n",
    "                padding=\"same\",\n",
    "                use_bias=False,\n",
    "            ),\n",
    "            # Reshape and Permute\n",
    "            layers.Reshape((-1, number_of_tokens)),  # (B, H*W, num_of_tokens)\n",
    "            layers.Permute((2, 1)),\n",
    "            layers.Lambda(lambda x: tf.nn.softmax(x, axis=-1)),\n",
    "        ]\n",
    "    )(\n",
    "        x\n",
    "    )  # (B, num_of_tokens, H*W)\n",
    "\n",
    "    # Reshape the input to align it with the output of the conv block.\n",
    "    num_filters = inputs.shape[-1]\n",
    "    inputs = layers.Conv2D(\n",
    "        filters=num_filters,\n",
    "        kernel_size=(1, 1),\n",
    "        strides=(1, 1),\n",
    "        groups=NUM_GROUPS,\n",
    "        padding=\"same\",\n",
    "        use_bias=False,\n",
    "    )(inputs)\n",
    "    inputs = layers.Reshape((-1, num_filters))(inputs)  # inputs == (B, H*W, C)\n",
    "\n",
    "    # Produced the attended inputs.\n",
    "    outputs = tf.einsum(\"...si,...id->...sd\", attention_maps, inputs)  # (B, num_tokens, C)\n",
    "    return outputs"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Transformer block"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def transformer(encoded_patches):\n",
    "    # Layer normalization 1.\n",
    "    x1 = layers.LayerNormalization(epsilon=LAYER_NORM_EPS)(encoded_patches)\n",
    "\n",
    "    # Multi Head Self Attention layer 1.\n",
    "    attention_output = layers.MultiHeadAttention(\n",
    "        num_heads=NUM_HEADS, key_dim=PROJECTION_DIM, dropout=0.1\n",
    "    )(x1, x1)\n",
    "\n",
    "    # Skip connection 1.\n",
    "    x2 = layers.Add()([attention_output, encoded_patches])\n",
    "\n",
    "    # Layer normalization 2.\n",
    "    x3 = layers.LayerNormalization(epsilon=LAYER_NORM_EPS)(x2)\n",
    "\n",
    "    # MLP layer 1.\n",
    "    x4 = mlp(x3, hidden_units=MLP_UNITS, dropout_rate=0.1)\n",
    "\n",
    "    # Skip connection 2.\n",
    "    encoded_patches = layers.Add()([x4, x2])\n",
    "    return encoded_patches"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "m3FSfbtFehYJ"
   },
   "source": [
    "## ViT model with optional TokenLearner"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "FKDtmPH-aNal"
   },
   "outputs": [],
   "source": [
    "def create_vit_classifier(use_token_learner=True, token_learner_units=NUM_TOKENS):\n",
    "    inputs = layers.Input(shape=INPUT_SHAPE)  # (B, H, W, C)\n",
    "\n",
    "    # Augment data.\n",
    "    augmented = data_augmentation(inputs)\n",
    "\n",
    "    # Create patches and project the pathces.\n",
    "    projected_patches = layers.Conv2D(\n",
    "        filters=PROJECTION_DIM,\n",
    "        kernel_size=(PATCH_SIZE, PATCH_SIZE),\n",
    "        strides=(PATCH_SIZE, PATCH_SIZE),\n",
    "        padding=\"VALID\",\n",
    "    )(augmented)\n",
    "    _, h, w, c = projected_patches.shape\n",
    "    projected_patches = layers.Reshape((h * w, c))(\n",
    "        projected_patches\n",
    "    )  # (B, number_patches, projection_dim)\n",
    "\n",
    "    # Add positional embeddings to the projected patches.\n",
    "    encoded_patches = position_embedding(\n",
    "        projected_patches\n",
    "    )  # (B, number_patches, projection_dim)\n",
    "    encoded_patches = layers.Dropout(0.1)(encoded_patches)\n",
    "\n",
    "    # Iterate over the number of layers and stack up blocks of\n",
    "    # Transformer.\n",
    "    for i in range(NUM_LAYERS):\n",
    "        # Add a Transformer block.\n",
    "        encoded_patches = transformer(encoded_patches)\n",
    "\n",
    "        # Add TokenLearner layer in the middle of the\n",
    "        # architecture. The paper suggests that anywhere\n",
    "        # between 1/2 or 3/4 will work well.\n",
    "        if use_token_learner and i == NUM_LAYERS // 2:\n",
    "            _, hh, c = encoded_patches.shape\n",
    "            h = int(math.sqrt(hh))\n",
    "            encoded_patches = layers.Reshape((h, h, c))(\n",
    "                encoded_patches\n",
    "            )  # (B, h, h, projection_dim)\n",
    "            encoded_patches = token_learner_v11(\n",
    "                encoded_patches, token_learner_units\n",
    "            )  # (B, num_tokens, c)\n",
    "\n",
    "    # Layer normalization and Global average pooling.\n",
    "    representation = layers.LayerNormalization(epsilon=LAYER_NORM_EPS)(encoded_patches)\n",
    "    representation = layers.GlobalAvgPool1D()(representation)\n",
    "\n",
    "    # Classify outputs.\n",
    "    outputs = layers.Dense(NUM_CLASSES, activation=\"softmax\")(representation)\n",
    "\n",
    "    # Create the Keras model.\n",
    "    model = keras.Model(inputs=inputs, outputs=outputs)\n",
    "    return model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "60oHNkcoef_5"
   },
   "source": [
    "## Training utility"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "sOmyzOhVfWoO"
   },
   "outputs": [],
   "source": [
    "def run_experiment(model, use_token_learner=True):\n",
    "    # Initialize the AdamW optimizer.\n",
    "    optimizer = tfa.optimizers.AdamW(\n",
    "        learning_rate=LEARNING_RATE, weight_decay=WEIGHT_DECAY\n",
    "    )\n",
    "    \n",
    "    # Compile the model with the optimizer, loss function\n",
    "    # and the metrics.\n",
    "    model.compile(\n",
    "        optimizer=optimizer,\n",
    "        loss=\"sparse_categorical_crossentropy\",\n",
    "        metrics=[\n",
    "            keras.metrics.SparseCategoricalAccuracy(name=\"accuracy\"),\n",
    "            keras.metrics.SparseTopKCategoricalAccuracy(5, name=\"top-5-accuracy\"),\n",
    "        ],\n",
    "    )\n",
    "    \n",
    "    # Define callbacks\n",
    "    checkpoint_filepath = \"/tmp/checkpoint\"\n",
    "    checkpoint_callback = keras.callbacks.ModelCheckpoint(\n",
    "        checkpoint_filepath,\n",
    "        monitor=\"val_accuracy\",\n",
    "        save_best_only=True,\n",
    "        save_weights_only=True,\n",
    "    )\n",
    "    timestamp = datetime.now().strftime(\"%Y%m%d%H%M%S\")\n",
    "    root_dir = \"logs-tokenlearner-v1.1\" if use_token_learner else \"logs-no-tokenlearner\"\n",
    "    tensorborad_callback = keras.callbacks.TensorBoard(\n",
    "        log_dir=f\"{root_dir}-{timestamp}\"\n",
    "    )\n",
    "    \n",
    "    # Train the model.\n",
    "    _ = model.fit(\n",
    "        train_ds,\n",
    "        epochs=EPOCHS,\n",
    "        validation_data=val_ds,\n",
    "        callbacks=[checkpoint_callback, tensorborad_callback],\n",
    "    )\n",
    "    \n",
    "    model.load_weights(checkpoint_filepath)\n",
    "    _, accuracy, top_5_accuracy = model.evaluate(test_ds)\n",
    "    print(f\"Test accuracy: {round(accuracy * 100, 2)}%\")\n",
    "    print(f\"Test top 5 accuracy: {round(top_5_accuracy * 100, 2)}%\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Rm2Q2b6hgY-s"
   },
   "source": [
    "## Experiments"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "pBWjky5s4d-m"
   },
   "outputs": [],
   "source": [
    "# Should at least be 5.\n",
    "num_trials = 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "_NMiWnSdgXXz"
   },
   "outputs": [],
   "source": [
    "# With TokenLearner.\n",
    "for _ in range(num_trials):\n",
    "    vit_token_learner = create_vit_classifier()\n",
    "    run_experiment(vit_token_learner, use_token_learner=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "BWlhIayN4d-o"
   },
   "source": [
    "## References\n",
    "\n",
    "* [Official TokenLearner code](https://github.com/google-research/scenic/blob/main/scenic/projects/token_learner/model.py)\n",
    "* [Image Classification with ViTs on keras.io](https://keras.io/examples/vision/image_classification_with_vision_transformer/)"
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "collapsed_sections": [],
   "machine_shape": "hm",
   "name": "TokenLearner.ipynb",
   "provenance": []
  },
  "environment": {
   "kernel": "python3",
   "name": "tf2-gpu.2-7.m84",
   "type": "gcloud",
   "uri": "gcr.io/deeplearning-platform-release/tf2-gpu.2-7:m84"
  },
  "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"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}