mae-pretraining.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "p-tCUmKj99CG"
   },
   "source": [
    "# Masked Autoencoders Are Scalable Vision Learners\n",
    "\n",
    "This notebook is a TF2.x implementation of [Masked Autoencoders Are Scalable Vision Learners](https://arxiv.org/abs/2111.06377) by He et. al.\n",
    "\n",
    "The notebook uses the following resource as a reference:\n",
    "\n",
    "- [Image classification with Vision Transformer](https://keras.io/examples/vision/image_classification_with_vision_transformer/)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "i8pvWN_--Veq",
    "tags": []
   },
   "source": [
    "# Import & Setup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "80ZKaTtG9zw9",
    "outputId": "467ea57f-03ac-43d5-e19a-a5a8ac99c393"
   },
   "outputs": [],
   "source": [
    "!pip install -q tensorflow-addons"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "NTZxVZLk-aLN"
   },
   "outputs": [],
   "source": [
    "from tensorflow.keras import layers\n",
    "import tensorflow_addons as tfa\n",
    "from tensorflow import keras\n",
    "import tensorflow as tf\n",
    "\n",
    "from datetime import datetime\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "import random\n",
    "\n",
    "# Setting seeds for reproducibility.\n",
    "SEED = 42\n",
    "tf.random.set_seed(SEED)\n",
    "np.random.seed(SEED)\n",
    "random.seed(SEED)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "oNM1SbqG-iHN"
   },
   "outputs": [],
   "source": [
    "# DATA\n",
    "BUFFER_SIZE = 1024\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 = 5e-3\n",
    "WEIGHT_DECAY = 1e-4\n",
    "\n",
    "# TRAINING\n",
    "EPOCHS = 100\n",
    "\n",
    "# AUGMENTATION\n",
    "IMAGE_SIZE = 48  # We'll resize input images to this size.\n",
    "PATCH_SIZE = 6  # Size of the patches to be extract from the input images.\n",
    "NUM_PATCHES = (IMAGE_SIZE // PATCH_SIZE) ** 2\n",
    "MASK_PROPORTION = 0.75\n",
    "\n",
    "# ENCODER and DECODER\n",
    "LAYER_NORM_EPS = 1e-6\n",
    "ENC_PROJECTION_DIM = 128\n",
    "DEC_PROJECTION_DIM = 64\n",
    "ENC_NUM_HEADS = 4\n",
    "ENC_LAYERS = 3\n",
    "DEC_NUM_HEADS = 4\n",
    "DEC_LAYERS = 1 # The decoder is lightweight but should be reasonably deep for reconstruction.\n",
    "ENC_TRANSFORMER_UNITS = [\n",
    "    ENC_PROJECTION_DIM * 2,\n",
    "    ENC_PROJECTION_DIM,\n",
    "]  # Size of the transformer layers.\n",
    "DEC_TRANSFORMER_UNITS = [\n",
    "    DEC_PROJECTION_DIM * 2,\n",
    "    DEC_PROJECTION_DIM,\n",
    "]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "r2qI0r9N_TeW"
   },
   "source": [
    "# Dataset\n",
    "\n",
    "Using **CIFAR10** for our implementation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "lMOYr_h1_QY6",
    "outputId": "e0201dc6-c933-4cf8-c54f-7fde8e3d27de"
   },
   "outputs": [],
   "source": [
    "(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)}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "sH3ihMqv_Zvn"
   },
   "outputs": [],
   "source": [
    "train_ds = tf.data.Dataset.from_tensor_slices((x_train))\n",
    "train_ds = train_ds.shuffle(BUFFER_SIZE).batch(BATCH_SIZE).prefetch(AUTO)\n",
    "\n",
    "val_ds = tf.data.Dataset.from_tensor_slices((x_val))\n",
    "val_ds = val_ds.batch(BATCH_SIZE).prefetch(AUTO)\n",
    "\n",
    "test_ds = tf.data.Dataset.from_tensor_slices((x_test))\n",
    "test_ds = test_ds.batch(BATCH_SIZE).prefetch(AUTO)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "I3jSERWD_dxh"
   },
   "source": [
    "# Data Augmentation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "6E12LNFc_dm5"
   },
   "outputs": [],
   "source": [
    "def get_train_augmentation_model():\n",
    "    model = 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=\"train_data_augmentation\",\n",
    "    )\n",
    "    return model\n",
    "\n",
    "\n",
    "def get_test_augmentation_model():\n",
    "    model = keras.Sequential(\n",
    "        [layers.Rescaling(1 / 255.0), layers.Resizing(IMAGE_SIZE, IMAGE_SIZE),],\n",
    "        name=\"test_data_augmentation\",\n",
    "    )\n",
    "    return model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "dynDm4z6AK-q"
   },
   "source": [
    "# Create Patches\n",
    "\n",
    "This layer creates patches from input images. The layer also consists of two utility methods:\n",
    "- `show_patched_image`: This utility function takes a batch of images and its corresponding patches, randomly choses a pair and plots it. This is useful for a sanity check.\n",
    "- `reconstruct_from_patch`: This utility funciton takes the patches of a **single** image, and reconstructs it back into the original image. This is useful for the training monitor callback defined later."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "dAiN1D-5AMGg"
   },
   "outputs": [],
   "source": [
    "class Patches(layers.Layer):\n",
    "    def __init__(self, patch_size=PATCH_SIZE, **kwargs):\n",
    "        super().__init__(**kwargs)\n",
    "        self.patch_size = patch_size\n",
    "\n",
    "        # Assuming the image has three channels each patch would be\n",
    "        # of size (patch_size, patch_size, 3).\n",
    "        self.resize = layers.Reshape((-1, patch_size * patch_size * 3))\n",
    "\n",
    "    def call(self, images):\n",
    "        # Create patches from the input images\n",
    "        patches = tf.image.extract_patches(\n",
    "            images=images,\n",
    "            sizes=[1, self.patch_size, self.patch_size, 1],\n",
    "            strides=[1, self.patch_size, self.patch_size, 1],\n",
    "            rates=[1, 1, 1, 1],\n",
    "            padding=\"VALID\",\n",
    "        )\n",
    "\n",
    "        # Reshape the patches to (batch, num_patches, patch_area) and return it.\n",
    "        patches = self.resize(patches)\n",
    "        return patches\n",
    "\n",
    "    def show_patched_image(self, images, patches):\n",
    "        # This is a utility function which accepts a batch of images and its\n",
    "        # corresponding patches and help visualize one image and its patches\n",
    "        # side by side.\n",
    "        idx = np.random.choice(patches.shape[0])\n",
    "        print(f\"Index selected: {idx}.\")\n",
    "\n",
    "        plt.figure(figsize=(4, 4))\n",
    "        plt.imshow(keras.utils.array_to_img(images[idx]))\n",
    "        plt.axis(\"off\")\n",
    "        plt.show()\n",
    "\n",
    "        n = int(np.sqrt(patches.shape[1]))\n",
    "        plt.figure(figsize=(4, 4))\n",
    "        for i, patch in enumerate(patches[idx]):\n",
    "            ax = plt.subplot(n, n, i + 1)\n",
    "            patch_img = tf.reshape(patch, (self.patch_size, self.patch_size, 3))\n",
    "            plt.imshow(keras.utils.img_to_array(patch_img))\n",
    "            plt.axis(\"off\")\n",
    "        plt.show()\n",
    "\n",
    "        # Return the index chosen to validate it outside the method.\n",
    "        return idx\n",
    "\n",
    "    # taken from https://stackoverflow.com/a/58082878/10319735\n",
    "    def reconstruct_from_patch(self, patch):\n",
    "        # This utility function takes patches from a *single* image and\n",
    "        # reconstructs it back into the image. This is useful for the train\n",
    "        # monitor callback.\n",
    "        num_patches = patch.shape[0]\n",
    "        n = int(np.sqrt(num_patches))\n",
    "        patch = tf.reshape(patch, (num_patches, self.patch_size, self.patch_size, 3))\n",
    "        rows = tf.split(patch, n, axis=0)\n",
    "        rows = [tf.concat(tf.unstack(x), axis=1) for x in rows]\n",
    "        reconstructed = tf.concat(rows, axis=0)\n",
    "        return reconstructed"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 496
    },
    "id": "ptI3I2aMB_rS",
    "outputId": "895d0714-f9ba-4235-ca8a-b3c21ec7091c"
   },
   "outputs": [],
   "source": [
    "# Get a batch of images.\n",
    "image_batch = next(iter(train_ds))\n",
    "\n",
    "# Augment the images.\n",
    "augmentation_model = get_train_augmentation_model()\n",
    "augmeneted_images = augmentation_model(image_batch)\n",
    "\n",
    "# Define the patch layer.\n",
    "patch_layer = Patches()\n",
    "\n",
    "# Get the patches from the batched images.\n",
    "patches = patch_layer(images=augmeneted_images)\n",
    "\n",
    "# Now pass the images and the corresponding patches\n",
    "# to the `show_patched_image` method.\n",
    "random_index = patch_layer.show_patched_image(images=augmeneted_images, patches=patches)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 248
    },
    "id": "qv0Va_68CPmF",
    "outputId": "5dd6d051-2740-452b-f57c-2131b7b8a3b1"
   },
   "outputs": [],
   "source": [
    "# Chose the same chose image and try reconstructing the patches\n",
    "# into the original image.\n",
    "image = patch_layer.reconstruct_from_patch(patches[random_index])\n",
    "plt.imshow(image)\n",
    "plt.axis(\"off\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "8GYmyvRZCbL4"
   },
   "source": [
    "# Patch Encoder\n",
    "\n",
    "This layer deals with encoding the pathces and adding the positional embedding too. The layer holds two utility functions:\n",
    "- `get_random_indices`: This function provides randomly sampled mask and unmask indices.\n",
    "- `show_masked_image`: A utility function that plots a random masked image."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "39WbLUN6CaU9"
   },
   "outputs": [],
   "source": [
    "class PatchEncoder(layers.Layer):\n",
    "    def __init__(\n",
    "        self,\n",
    "        patch_size=PATCH_SIZE,\n",
    "        projection_dim=ENC_PROJECTION_DIM,\n",
    "        mask_proportion=MASK_PROPORTION,\n",
    "        downstream=False,\n",
    "        **kwargs,\n",
    "    ):\n",
    "        super().__init__(**kwargs)\n",
    "        self.patch_size = patch_size\n",
    "        self.projection_dim = projection_dim\n",
    "        self.mask_proportion = mask_proportion\n",
    "        self.downstream = downstream\n",
    "\n",
    "        # This is a trainable mask token initialized randomly from a normal\n",
    "        # distribution.\n",
    "        self.mask_token = tf.Variable(\n",
    "            tf.random.normal([1, patch_size * patch_size * 3]), trainable=True\n",
    "        )\n",
    "\n",
    "    def build(self, input_shape):\n",
    "        (_, self.num_patches, self.patch_area) = input_shape\n",
    "\n",
    "        # Create the projection layer for the patches.\n",
    "        self.projection = layers.Dense(units=self.projection_dim)\n",
    "\n",
    "        # Create the positional embedding layer.\n",
    "        self.position_embedding = layers.Embedding(\n",
    "            input_dim=self.num_patches, output_dim=self.projection_dim\n",
    "        )\n",
    "\n",
    "        # Number of patches that will be masked.\n",
    "        self.num_mask = int(self.mask_proportion * self.num_patches)\n",
    "\n",
    "    def call(self, patches):\n",
    "        # Get the positional embeddings.\n",
    "        batch_size = tf.shape(patches)[0]\n",
    "        positions = tf.range(start=0, limit=self.num_patches, delta=1)\n",
    "        pos_embeddings = self.position_embedding(positions[tf.newaxis, ...])\n",
    "        pos_embeddings = tf.tile(\n",
    "            pos_embeddings, [batch_size, 1, 1]\n",
    "        )  # (B, num_patches, projection_dim)\n",
    "\n",
    "        # Embed the patches.\n",
    "        patch_embeddings = (\n",
    "            self.projection(patches) + pos_embeddings\n",
    "        )  # (B, num_patches, projection_dim)\n",
    "\n",
    "        if self.downstream:\n",
    "            return patch_embeddings\n",
    "        else:\n",
    "            mask_indices, unmask_indices = self.get_random_indices(batch_size)\n",
    "            # The encoder input is the unmasked patch embeddings. Here we gather\n",
    "            # all the patches that should be unmasked.\n",
    "            unmasked_embeddings = tf.gather(\n",
    "                patch_embeddings, unmask_indices, axis=1, batch_dims=1\n",
    "            )  # (B, unmask_numbers, projection_dim)\n",
    "\n",
    "            # Get the unmasked and masked position embeddings. We will need them\n",
    "            # for the decoder.\n",
    "            unmasked_positions = tf.gather(\n",
    "                pos_embeddings, unmask_indices, axis=1, batch_dims=1\n",
    "            )  # (B, unmask_numbers, projection_dim)\n",
    "            masked_positions = tf.gather(\n",
    "                pos_embeddings, mask_indices, axis=1, batch_dims=1\n",
    "            )  # (B, mask_numbers, projection_dim)\n",
    "\n",
    "            # Repeat the mask token number of mask times.\n",
    "            # Mask tokens replace the masks of the image.\n",
    "            mask_tokens = tf.repeat(self.mask_token, repeats=self.num_mask, axis=0)\n",
    "            mask_tokens = tf.repeat(\n",
    "                mask_tokens[tf.newaxis, ...], repeats=batch_size, axis=0\n",
    "            )\n",
    "\n",
    "            # Get the masked embeddings for the tokens.\n",
    "            masked_embeddings = self.projection(mask_tokens) + masked_positions\n",
    "            return (\n",
    "                unmasked_embeddings,  # input to the encoder\n",
    "                masked_embeddings,  # first part of input to the decoder\n",
    "                unmasked_positions,  # added to the encoder outputs\n",
    "                mask_indices,  # the indices that were masked\n",
    "                unmask_indices,  # the indices that were unmaksed\n",
    "            )\n",
    "\n",
    "    def get_random_indices(self, batch_size):\n",
    "        # Create random indices from a uniform distribution and then split\n",
    "        # it into mask and unmask indices.\n",
    "        rand_indices = tf.argsort(\n",
    "            tf.random.uniform(shape=(batch_size, self.num_patches)), axis=-1\n",
    "        )\n",
    "        mask_indices = rand_indices[:, : self.num_mask]\n",
    "        unmask_indices = rand_indices[:, self.num_mask :]\n",
    "\n",
    "        return mask_indices, unmask_indices\n",
    "\n",
    "    def show_masked_image(self, patches, unmask_indices):\n",
    "        # choose a random patch and it corresponding unmask index\n",
    "        idx = np.random.choice(patches.shape[0])\n",
    "        patch = patches[idx]\n",
    "        unmask_index = unmask_indices[idx]\n",
    "\n",
    "        # build a numpy array of same shape as pathc\n",
    "        new_patch = np.zeros_like(patch)\n",
    "\n",
    "        # iterate of the new_patch and plug the unmasked patches\n",
    "        count = 0\n",
    "        for i in range(unmask_index.shape[0]):\n",
    "            new_patch[unmask_index[i]] = patch[unmask_index[i]]\n",
    "        return new_patch, idx"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "PvEhNFDcFCLF"
   },
   "outputs": [],
   "source": [
    "# Create the patch encoder layer.\n",
    "patch_encoder = PatchEncoder()\n",
    "\n",
    "# Get the embeddings and positions.\n",
    "(\n",
    "    unmasked_embeddings,\n",
    "    masked_embeddings,\n",
    "    unmasked_positions,\n",
    "    mask_indices,\n",
    "    unmask_indices,\n",
    ") = patch_encoder(patches=patches)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 301
    },
    "id": "UlxangzdFwMJ",
    "outputId": "72caf44d-71d8-4e95-b66d-f4a7a75b0c90"
   },
   "outputs": [],
   "source": [
    "# Show a maksed patch image.\n",
    "new_patch, random_index = patch_encoder.show_masked_image(patches, unmask_indices)\n",
    "\n",
    "plt.figure(figsize=(10, 10))\n",
    "plt.subplot(1, 2, 1)\n",
    "img = patch_layer.reconstruct_from_patch(new_patch)\n",
    "plt.imshow(keras.utils.array_to_img(img))\n",
    "plt.axis(\"off\")\n",
    "plt.title(\"Masked\")\n",
    "plt.subplot(1, 2, 2)\n",
    "img = augmeneted_images[random_index]\n",
    "plt.imshow(keras.utils.array_to_img(img))\n",
    "plt.axis(\"off\")\n",
    "plt.title(\"Original\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "A88s55FzF9Rz"
   },
   "source": [
    "# MLP"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "91-B5HvSF52o"
   },
   "outputs": [],
   "source": [
    "def mlp(x, dropout_rate, hidden_units):\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": "A_CHs2O1GA8X"
   },
   "source": [
    "# Encoder"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "56BDy60iGB4N"
   },
   "outputs": [],
   "source": [
    "def create_encoder(num_heads=ENC_NUM_HEADS, num_layers=ENC_LAYERS):\n",
    "    inputs = layers.Input((None, ENC_PROJECTION_DIM))\n",
    "    x = inputs\n",
    "\n",
    "    for _ in range(num_layers):\n",
    "        # Layer normalization 1.\n",
    "        x1 = layers.LayerNormalization(epsilon=LAYER_NORM_EPS)(x)\n",
    "\n",
    "        # Create a multi-head attention layer.\n",
    "        attention_output = layers.MultiHeadAttention(\n",
    "            num_heads=num_heads, key_dim=ENC_PROJECTION_DIM, dropout=0.1\n",
    "        )(x1, x1)\n",
    "\n",
    "        # Skip connection 1.\n",
    "        x2 = layers.Add()([attention_output, x])\n",
    "\n",
    "        # Layer normalization 2.\n",
    "        x3 = layers.LayerNormalization(epsilon=LAYER_NORM_EPS)(x2)\n",
    "\n",
    "        # MLP.\n",
    "        x3 = mlp(x3, hidden_units=ENC_TRANSFORMER_UNITS, dropout_rate=0.1)\n",
    "\n",
    "        # Skip connection 2.\n",
    "        x = layers.Add()([x3, x2])\n",
    "\n",
    "    outputs = layers.LayerNormalization(epsilon=LAYER_NORM_EPS)(x)\n",
    "    return keras.Model(inputs, outputs, name=\"mae_encoder\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "i-r40qxxGOnF"
   },
   "source": [
    "# Decoder"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "u4t3AnQ1GN8i"
   },
   "outputs": [],
   "source": [
    "def create_decoder(\n",
    "    num_layers=DEC_LAYERS, num_heads=DEC_NUM_HEADS, image_size=IMAGE_SIZE\n",
    "):\n",
    "    inputs = layers.Input((NUM_PATCHES, ENC_PROJECTION_DIM))\n",
    "    x = layers.Dense(DEC_PROJECTION_DIM)(inputs)\n",
    "\n",
    "    for _ in range(num_layers):\n",
    "        # Layer normalization 1.\n",
    "        x1 = layers.LayerNormalization(epsilon=LAYER_NORM_EPS)(x)\n",
    "\n",
    "        # Create a multi-head attention layer.\n",
    "        attention_output = layers.MultiHeadAttention(\n",
    "            num_heads=num_heads, key_dim=DEC_PROJECTION_DIM, dropout=0.1\n",
    "        )(x1, x1)\n",
    "\n",
    "        # Skip connection 1.\n",
    "        x2 = layers.Add()([attention_output, x])\n",
    "\n",
    "        # Layer normalization 2.\n",
    "        x3 = layers.LayerNormalization(epsilon=LAYER_NORM_EPS)(x2)\n",
    "\n",
    "        # MLP.\n",
    "        x3 = mlp(x3, hidden_units=DEC_TRANSFORMER_UNITS, dropout_rate=0.1)\n",
    "\n",
    "        # Skip connection 2.\n",
    "        x = layers.Add()([x3, x2])\n",
    "\n",
    "    x = layers.LayerNormalization(epsilon=LAYER_NORM_EPS)(x)\n",
    "    x = layers.Flatten()(x)\n",
    "    pre_final = layers.Dense(units=image_size * image_size * 3, activation=\"sigmoid\")(x)\n",
    "    outputs = layers.Reshape((image_size, image_size, 3))(pre_final)\n",
    "\n",
    "    return keras.Model(inputs, outputs, name=\"mae_decoder\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "KQ1LwTdlGcuU"
   },
   "source": [
    "# MaskedAutoEncoder Model\n",
    "\n",
    "This is the trainer model where we encapsulate the training logic."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "7-V6L61yGbg6"
   },
   "outputs": [],
   "source": [
    "class MaskedAutoencoder(keras.Model):\n",
    "    def __init__(\n",
    "        self,\n",
    "        train_augmentation_model,\n",
    "        test_augmentation_model,\n",
    "        patch_layer,\n",
    "        patch_encoder,\n",
    "        encoder,\n",
    "        decoder,\n",
    "        **kwargs\n",
    "    ):\n",
    "        super().__init__(**kwargs)\n",
    "        self.train_augmentation_model = train_augmentation_model\n",
    "        self.test_augmentation_model = test_augmentation_model\n",
    "        self.patch_layer = patch_layer\n",
    "        self.patch_encoder = patch_encoder\n",
    "        self.encoder = encoder\n",
    "        self.decoder = decoder\n",
    "\n",
    "    def calculate_loss(self, images, test=False):\n",
    "        # Augment the input images.\n",
    "        if test:\n",
    "            augmeneted_images = self.test_augmentation_model(images)\n",
    "        else:\n",
    "            augmeneted_images = self.train_augmentation_model(images)\n",
    "\n",
    "        # Patch the augmented images.\n",
    "        patches = self.patch_layer(augmeneted_images)\n",
    "\n",
    "        # Encode the patches.\n",
    "        (\n",
    "            unmasked_embeddings,\n",
    "            masked_embeddings,\n",
    "            unmasked_positions,\n",
    "            mask_indices,\n",
    "            unmask_indices,\n",
    "        ) = self.patch_encoder(patches)\n",
    "\n",
    "        # Pass the unmaksed patche to the encoder.\n",
    "        encoder_outputs = self.encoder(unmasked_embeddings)\n",
    "\n",
    "        # Create the decoder inputs.\n",
    "        encoder_outputs = encoder_outputs + unmasked_positions\n",
    "        decoder_inputs = tf.concat([encoder_outputs, masked_embeddings], axis=1)\n",
    "\n",
    "        # Decode the inputs.\n",
    "        decoder_outputs = self.decoder(decoder_inputs)\n",
    "        decoder_patches = self.patch_layer(decoder_outputs)\n",
    "\n",
    "        loss_patch = tf.gather(patches, mask_indices, axis=1, batch_dims=1)\n",
    "        loss_output = tf.gather(decoder_patches, mask_indices, axis=1, batch_dims=1)\n",
    "\n",
    "        # Compute the total loss.\n",
    "        total_loss = self.compiled_loss(loss_patch, loss_output)\n",
    "\n",
    "        return total_loss, loss_patch, loss_output\n",
    "\n",
    "    def train_step(self, images):\n",
    "        with tf.GradientTape() as tape:\n",
    "            total_loss, loss_patch, loss_output = self.calculate_loss(images)\n",
    "\n",
    "        # Apply gradients.\n",
    "        train_vars = [\n",
    "            self.train_augmentation_model.trainable_variables,\n",
    "            self.patch_layer.trainable_variables,\n",
    "            self.patch_encoder.trainable_variables,\n",
    "            self.encoder.trainable_variables,\n",
    "            self.decoder.trainable_variables,\n",
    "        ]\n",
    "        grads = tape.gradient(total_loss, train_vars)\n",
    "        tv_list = []\n",
    "        for (grad, var) in zip(grads, train_vars):\n",
    "            for g, v in zip(grad, var):\n",
    "                tv_list.append((g, v))\n",
    "        self.optimizer.apply_gradients(tv_list)\n",
    "\n",
    "        # Report progress.\n",
    "        self.compiled_metrics.update_state(loss_patch, loss_output)\n",
    "        return {m.name: m.result() for m in self.metrics}\n",
    "\n",
    "    def test_step(self, images):\n",
    "        total_loss, loss_patch, loss_output = self.calculate_loss(images, test=True)\n",
    "\n",
    "        # Update the trackers.\n",
    "        self.compiled_metrics.update_state(loss_patch, loss_output)\n",
    "        return {m.name: m.result() for m in self.metrics}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "oCOI8_9BX_6g"
   },
   "source": [
    "# Model init"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "2LqZWL1eHS8l"
   },
   "outputs": [],
   "source": [
    "keras.backend.clear_session()\n",
    "\n",
    "train_augmentation_model = get_train_augmentation_model()\n",
    "test_augmentation_model = get_test_augmentation_model()\n",
    "patch_layer = Patches()\n",
    "patch_encoder = PatchEncoder()\n",
    "encoder = create_encoder()\n",
    "decoder = create_decoder()\n",
    "\n",
    "mae_model = MaskedAutoencoder(\n",
    "    train_augmentation_model=train_augmentation_model,\n",
    "    test_augmentation_model=test_augmentation_model,\n",
    "    patch_layer=patch_layer,\n",
    "    patch_encoder=patch_encoder,\n",
    "    encoder=encoder,\n",
    "    decoder=decoder,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "TCDOA_kdX_6h"
   },
   "source": [
    "## Training callbacks"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "QjuUx4OYHaQx"
   },
   "outputs": [],
   "source": [
    "# Taking a batch of test inputs to measure model's progress.\n",
    "test_images = next(iter(test_ds))\n",
    "\n",
    "\n",
    "class TrainMonitor(tf.keras.callbacks.Callback):\n",
    "    def __init__(self, epoch_interval=None):\n",
    "        self.epoch_interval = epoch_interval\n",
    "\n",
    "    def on_epoch_end(self, epoch, logs=None):\n",
    "        if self.epoch_interval and epoch % self.epoch_interval == 0:\n",
    "            test_augmeneted_images = self.model.test_augmentation_model(test_images)\n",
    "            test_patches = self.model.patch_layer(test_augmeneted_images)\n",
    "            (\n",
    "                test_unmasked_embeddings,\n",
    "                test_masked_embeddings,\n",
    "                test_unmasked_positions,\n",
    "                test_mask_indices,\n",
    "                test_unmask_indices,\n",
    "            ) = self.model.patch_encoder(test_patches)\n",
    "            test_encoder_outputs = self.model.encoder(test_unmasked_embeddings)\n",
    "            test_encoder_outputs = test_encoder_outputs + test_unmasked_positions\n",
    "            test_decoder_inputs = tf.concat(\n",
    "                [test_encoder_outputs, test_masked_embeddings], axis=1\n",
    "            )\n",
    "            test_decoder_outputs = self.model.decoder(test_decoder_inputs)\n",
    "\n",
    "            # Show a maksed patch image.\n",
    "            test_masked_patch, idx = self.model.patch_encoder.show_masked_image(\n",
    "                test_patches, test_unmask_indices\n",
    "            )\n",
    "            print(f\"\\nIdx chosen: {idx}\")\n",
    "            original_image = test_augmeneted_images[idx]\n",
    "            masked_image = self.model.patch_layer.reconstruct_from_patch(\n",
    "                test_masked_patch\n",
    "            )\n",
    "            reconstructed_image = test_decoder_outputs[idx]\n",
    "\n",
    "            fig, ax = plt.subplots(nrows=1, ncols=3, figsize=(15, 5))\n",
    "            ax[0].imshow(original_image)\n",
    "            ax[0].set_title(f\"Original: {epoch:03d}\")\n",
    "\n",
    "            ax[1].imshow(masked_image)\n",
    "            ax[1].set_title(f\"Masked: {epoch:03d}\")\n",
    "\n",
    "            ax[2].imshow(reconstructed_image)\n",
    "            ax[2].set_title(f\"Resonstructed: {epoch:03d}\")\n",
    "\n",
    "            plt.show()\n",
    "            plt.close()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "s4PEJxppX_6h"
   },
   "outputs": [],
   "source": [
    "# Some code is taken from:\n",
    "# https://www.kaggle.com/ashusma/training-rfcx-tensorflow-tpu-effnet-b2.\n",
    "\n",
    "\n",
    "class WarmUpCosine(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",
    "\n",
    "        cos_annealed_lr = tf.cos(\n",
    "            self.pi\n",
    "            * (tf.cast(step, tf.float32) - self.warmup_steps)\n",
    "            / float(self.total_steps - self.warmup_steps)\n",
    "        )\n",
    "        learning_rate = 0.5 * self.learning_rate_base * (1 + cos_annealed_lr)\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",
    "        )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 283
    },
    "id": "iTn6VcaBX_6h",
    "outputId": "c0df7a58-763f-4a1f-bbf3-1e1f2e64784f"
   },
   "outputs": [],
   "source": [
    "total_steps = int((len(x_train) / BATCH_SIZE) * EPOCHS)\n",
    "warmup_steps = int(total_steps * 0.15)\n",
    "scheduled_lrs = WarmUpCosine(\n",
    "    learning_rate_base=LEARNING_RATE,\n",
    "    total_steps=total_steps,\n",
    "    warmup_learning_rate=0.0,\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()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "l_heFaYxJaOp"
   },
   "outputs": [],
   "source": [
    "timestamp = datetime.utcnow().strftime(\"%y%m%d-%H%M%S\")\n",
    "\n",
    "train_callbacks = [\n",
    "    keras.callbacks.TensorBoard(log_dir=f\"mae_logs_{timestamp}\"),\n",
    "    TrainMonitor(epoch_interval=5),\n",
    "]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "-dUZ1RYqX_6h"
   },
   "source": [
    "# Compilation and training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "gtgH-OxyJczw"
   },
   "outputs": [],
   "source": [
    "optimizer = tfa.optimizers.AdamW(learning_rate=scheduled_lrs, weight_decay=WEIGHT_DECAY)\n",
    "\n",
    "mae_model.compile(\n",
    "    optimizer=optimizer, loss=keras.losses.MeanSquaredError(), metrics=[\"mae\"]\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 1000
    },
    "id": "ZUAXzpDoJiXG",
    "outputId": "2015187b-0bfb-445e-a687-018a9284a649"
   },
   "outputs": [],
   "source": [
    "history = mae_model.fit(\n",
    "    train_ds, epochs=EPOCHS, validation_data=val_ds, callbacks=train_callbacks,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "S0aVUm63Lj-L",
    "outputId": "cbe80727-6cc3-4db3-ec5d-cf01469f38d0"
   },
   "outputs": [],
   "source": [
    "loss, mae = mae_model.evaluate(test_ds)\n",
    "print(f\"Loss: {loss:.2f}\")\n",
    "print(f\"MAE: {mae:.2f}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "u8mJ1QkHLoPd"
   },
   "source": [
    "# Downstrean Task"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "kXnO5jNALndF",
    "outputId": "8b6c45a7-10bf-4e58-f866-f691b5352b12"
   },
   "outputs": [],
   "source": [
    "# Extract the augmentation layers.\n",
    "train_augmentation_model = mae_model.train_augmentation_model\n",
    "test_augmentation_model = mae_model.test_augmentation_model\n",
    "\n",
    "# Extract the patchers.\n",
    "patch_layer = mae_model.patch_layer\n",
    "patch_encoder = mae_model.patch_encoder\n",
    "patch_encoder.downstream = True  # Swtich the downstream flag to True.\n",
    "\n",
    "# Extract the encoder.\n",
    "encoder = mae_model.encoder\n",
    "\n",
    "# Pack as a model.\n",
    "downstream_model = keras.Sequential(\n",
    "    [\n",
    "        layers.Input((IMAGE_SIZE, IMAGE_SIZE, 3)),\n",
    "        patch_layer,\n",
    "        patch_encoder,\n",
    "        encoder,\n",
    "        layers.BatchNormalization(),  # Refer to A.1 (Linear probing)\n",
    "        layers.GlobalAveragePooling1D(),\n",
    "        layers.Dense(NUM_CLASSES, activation=\"softmax\"),\n",
    "    ],\n",
    "    name=\"linear_probe_model\",\n",
    ")\n",
    "\n",
    "# Only the final classification layer of the `downstream_model` should be trainable.\n",
    "for layer in downstream_model.layers[:-1]:\n",
    "    layer.trainable = False\n",
    "\n",
    "downstream_model.summary()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "oPaNGdfvLtyR"
   },
   "outputs": [],
   "source": [
    "def prepare_data(images, labels, is_train=True):\n",
    "    if is_train:\n",
    "        augmentation_model = train_augmentation_model\n",
    "    else:\n",
    "        augmentation_model = test_augmentation_model\n",
    "\n",
    "    dataset = tf.data.Dataset.from_tensor_slices((images, labels))\n",
    "    if is_train:\n",
    "        dataset = dataset.shuffle(BUFFER_SIZE)\n",
    "\n",
    "    dataset = dataset.batch(BATCH_SIZE).map(\n",
    "        lambda x, y: (augmentation_model(x), y), num_parallel_calls=AUTO\n",
    "    )\n",
    "    return dataset.prefetch(AUTO)\n",
    "\n",
    "\n",
    "train_ds = prepare_data(x_train, y_train)\n",
    "val_ds = prepare_data(x_train, y_train, is_train=False)\n",
    "test_ds = prepare_data(x_test, y_test, is_train=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "xdeuZ98oLvis",
    "outputId": "c9f7b566-47e9-4afa-a86d-728df915cd99"
   },
   "outputs": [],
   "source": [
    "linear_probe_epochs = 50\n",
    "linear_prob_lr = 0.1\n",
    "warm_epoch_percentage = 0.1\n",
    "steps = int((len(x_train) // BATCH_SIZE) * linear_probe_epochs)\n",
    "\n",
    "warmup_steps = int(steps * warm_epoch_percentage)\n",
    "scheduled_lrs = WarmUpCosine(\n",
    "    learning_rate_base=linear_prob_lr,\n",
    "    total_steps=steps,\n",
    "    warmup_learning_rate=0.0,\n",
    "    warmup_steps=warmup_steps,\n",
    ")\n",
    "\n",
    "optimizer = keras.optimizers.SGD(learning_rate=scheduled_lrs, momentum=0.9)\n",
    "downstream_model.compile(\n",
    "    optimizer=optimizer, loss=\"sparse_categorical_crossentropy\", metrics=[\"accuracy\"]\n",
    ")\n",
    "downstream_model.fit(train_ds, validation_data=val_ds, epochs=linear_probe_epochs)\n",
    "\n",
    "loss, accuracy = downstream_model.evaluate(test_ds)\n",
    "accuracy = round(accuracy * 100, 2)\n",
    "print(f\"Accuracy on the test set: {accuracy}%.\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "0o8DHVfAMrSL",
    "outputId": "c98b307d-c002-4507-ac21-0ca777cf3f06"
   },
   "outputs": [],
   "source": [
    "downstream_model.save(f\"linear_probe_{timestamp}\", include_optimizer=False)"
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "collapsed_sections": [],
   "machine_shape": "hm",
   "name": "mae-pretraining.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",
   "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.9.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}