SimCLR_ImageNet_Subset.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "name": "SimCLR_ImageNet_Subset.ipynb",
      "provenance": [],
      "toc_visible": true,
      "include_colab_link": true
    },
    "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.7.1"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/sayakpaul/SimCLR-in-TensorFlow-2/blob/master/SimCLR_ImageNet_Subset.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "LMp5lKNjgBXq",
        "colab_type": "text"
      },
      "source": [
        "## Initial setup"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab_type": "code",
        "id": "lIYdn1woOS1n",
        "colab": {}
      },
      "source": [
        "import tensorflow as tf\n",
        "print(tf.__version__)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab_type": "code",
        "id": "dZowsDvOYK37",
        "colab": {}
      },
      "source": [
        "!nvidia-smi"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "I2MOoF9zgMHp",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# Install wandb for experiment tracking\n",
        "!pip install --upgrade https://github.com/wandb/client/archive/feature/code-save.zip"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "EXx0RnrigMwd",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "import wandb\n",
        "wandb.login()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab_type": "code",
        "id": "421MrJNMYQD7",
        "colab": {}
      },
      "source": [
        "# Other imports\n",
        "from tensorflow.keras.layers import *\n",
        "from tensorflow.keras.models import *\n",
        "import matplotlib.pyplot as plt\n",
        "from imutils import paths\n",
        "from tqdm import tqdm\n",
        "import tensorflow as tf\n",
        "import numpy as np\n",
        "import cv2\n",
        "\n",
        "# Random seed fixation\n",
        "tf.random.set_seed(666)\n",
        "np.random.seed(666)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "UUFlGxuJgBX9",
        "colab_type": "text"
      },
      "source": [
        "## Dataset gathering and preparation"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab_type": "code",
        "id": "cq8xwssDYNTq",
        "colab": {}
      },
      "source": [
        "# Gather dataset\n",
        "!git clone https://github.com/thunderInfy/imagenet-5-categories"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab_type": "code",
        "id": "Z0atoQ0MYT4b",
        "colab": {}
      },
      "source": [
        "# How many training images for SimCLR?\n",
        "train_images = list(paths.list_images(\"imagenet-5-categories/train\"))\n",
        "print(len(train_images))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab_type": "code",
        "id": "Ad9dX7ONYV9Y",
        "colab": {}
      },
      "source": [
        "# Augmentation utilities (differs from the original implementation)\n",
        "# Referred from: https://arxiv.org/pdf/2002.05709.pdf (Appendxi A \n",
        "# corresponding GitHub: https://github.com/google-research/simclr/)\n",
        "\n",
        "class CustomAugment(object):\n",
        "    def __call__(self, sample):        \n",
        "        # Random flips\n",
        "        sample = self._random_apply(tf.image.flip_left_right, sample, p=0.5)\n",
        "        \n",
        "        # Randomly apply transformation (color distortions) with probability p.\n",
        "        sample = self._random_apply(self._color_jitter, sample, p=0.8)\n",
        "        sample = self._random_apply(self._color_drop, sample, p=0.2)\n",
        "\n",
        "        return sample\n",
        "\n",
        "    def _color_jitter(self, x, s=1):\n",
        "        # one can also shuffle the order of following augmentations\n",
        "        # each time they are applied.\n",
        "        x = tf.image.random_brightness(x, max_delta=0.8*s)\n",
        "        x = tf.image.random_contrast(x, lower=1-0.8*s, upper=1+0.8*s)\n",
        "        x = tf.image.random_saturation(x, lower=1-0.8*s, upper=1+0.8*s)\n",
        "        x = tf.image.random_hue(x, max_delta=0.2*s)\n",
        "        x = tf.clip_by_value(x, 0, 1)\n",
        "        return x\n",
        "    \n",
        "    def _color_drop(self, x):\n",
        "        x = tf.image.rgb_to_grayscale(x)\n",
        "        x = tf.tile(x, [1, 1, 1, 3])\n",
        "        return x\n",
        "    \n",
        "    def _random_apply(self, func, x, p):\n",
        "        return tf.cond(\n",
        "          tf.less(tf.random.uniform([], minval=0, maxval=1, dtype=tf.float32),\n",
        "                  tf.cast(p, tf.float32)),\n",
        "          lambda: func(x),\n",
        "          lambda: x)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab_type": "code",
        "id": "XKYH6UxnYa7g",
        "colab": {}
      },
      "source": [
        "# Build the augmentation pipeline\n",
        "data_augmentation = Sequential([Lambda(CustomAugment())])"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab_type": "code",
        "id": "4eDTVhuHYdwi",
        "colab": {}
      },
      "source": [
        "# Image preprocessing utils\n",
        "@tf.function\n",
        "def parse_images(image_path):\n",
        "    image_string = tf.io.read_file(image_path)\n",
        "    image = tf.image.decode_jpeg(image_string, channels=3)\n",
        "    image = tf.image.convert_image_dtype(image, tf.float32)\n",
        "    image = tf.image.resize(image, size=[224, 224])\n",
        "\n",
        "    return image"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab_type": "code",
        "id": "otqf-I5FZLqd",
        "colab": {}
      },
      "source": [
        "# Create TensorFlow dataset\n",
        "BATCH_SIZE = 64\n",
        "\n",
        "train_ds = tf.data.Dataset.from_tensor_slices(train_images)\n",
        "train_ds = (\n",
        "    train_ds\n",
        "    .map(parse_images, num_parallel_calls=tf.data.experimental.AUTOTUNE)\n",
        "    .shuffle(1024)\n",
        "    .batch(BATCH_SIZE, drop_remainder=True)\n",
        "    .prefetch(tf.data.experimental.AUTOTUNE)\n",
        ")"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "2avjHJzRgBYS",
        "colab_type": "text"
      },
      "source": [
        "## Utilities"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab_type": "code",
        "id": "F1_Wl0DzZQ5F",
        "colab": {}
      },
      "source": [
        "# Architecture utils\n",
        "def get_resnet_simclr(hidden_1, hidden_2, hidden_3):\n",
        "    base_model = tf.keras.applications.ResNet50(include_top=False, weights=None, input_shape=(224, 224, 3))\n",
        "    base_model.trainable = True\n",
        "    inputs = Input((224, 224, 3))\n",
        "    h = base_model(inputs, training=True)\n",
        "    h = GlobalAveragePooling2D()(h)\n",
        "\n",
        "    projection_1 = Dense(hidden_1)(h)\n",
        "    projection_1 = Activation(\"relu\")(projection_1)\n",
        "    projection_2 = Dense(hidden_2)(projection_1)\n",
        "    projection_2 = Activation(\"relu\")(projection_2)\n",
        "    projection_3 = Dense(hidden_3)(projection_2)\n",
        "\n",
        "    resnet_simclr = Model(inputs, projection_3)\n",
        "\n",
        "    return resnet_simclr"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab_type": "code",
        "id": "7-hwmH_LZc7K",
        "colab": {}
      },
      "source": [
        "!wget https://raw.githubusercontent.com/sthalles/SimCLR-tensorflow/master/utils/helpers.py\n",
        "!wget https://raw.githubusercontent.com/sthalles/SimCLR-tensorflow/master/utils/losses.py"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "X6d1jsrs7ztp"
      },
      "source": [
        "Don't forget to comment the `augmentation` import in the `helpers.py` script."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab_type": "code",
        "id": "2OUhSNfTZgY6",
        "colab": {}
      },
      "source": [
        "from losses import _dot_simililarity_dim1 as sim_func_dim1, _dot_simililarity_dim2 as sim_func_dim2\n",
        "import helpers"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab_type": "code",
        "id": "Y0Ae5dMLZney",
        "colab": {}
      },
      "source": [
        "# Mask to remove positive examples from the batch of negative samples\n",
        "negative_mask = helpers.get_negative_mask(BATCH_SIZE)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab_type": "code",
        "id": "K6qPdhPrZpFA",
        "colab": {}
      },
      "source": [
        "@tf.function\n",
        "def train_step(xis, xjs, model, optimizer, criterion, temperature):\n",
        "    with tf.GradientTape() as tape:\n",
        "        zis = model(xis)\n",
        "        zjs = model(xjs)\n",
        "\n",
        "        # normalize projection feature vectors\n",
        "        zis = tf.math.l2_normalize(zis, axis=1)\n",
        "        zjs = tf.math.l2_normalize(zjs, axis=1)\n",
        "\n",
        "        l_pos = sim_func_dim1(zis, zjs)\n",
        "        l_pos = tf.reshape(l_pos, (BATCH_SIZE, 1))\n",
        "        l_pos /= temperature\n",
        "\n",
        "        negatives = tf.concat([zjs, zis], axis=0)\n",
        "\n",
        "        loss = 0\n",
        "\n",
        "        for positives in [zis, zjs]:\n",
        "            l_neg = sim_func_dim2(positives, negatives)\n",
        "\n",
        "            labels = tf.zeros(BATCH_SIZE, dtype=tf.int32)\n",
        "\n",
        "            l_neg = tf.boolean_mask(l_neg, negative_mask)\n",
        "            l_neg = tf.reshape(l_neg, (BATCH_SIZE, -1))\n",
        "            l_neg /= temperature\n",
        "\n",
        "            logits = tf.concat([l_pos, l_neg], axis=1) \n",
        "            loss += criterion(y_pred=logits, y_true=labels)\n",
        "\n",
        "        loss = loss / (2 * BATCH_SIZE)\n",
        "\n",
        "    gradients = tape.gradient(loss, model.trainable_variables)\n",
        "    optimizer.apply_gradients(zip(gradients, model.trainable_variables))\n",
        "\n",
        "    return loss"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "nSs5WYn7g5v_",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "wandb.init(project=\"simclr\")"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab_type": "code",
        "id": "zfp8gT4ZZtif",
        "colab": {}
      },
      "source": [
        "def train_simclr(model, dataset, optimizer, criterion,\n",
        "                 temperature=0.1, epochs=100):\n",
        "    step_wise_loss = []\n",
        "    epoch_wise_loss = []\n",
        "\n",
        "    for epoch in tqdm(range(epochs)):\n",
        "        for image_batch in dataset:\n",
        "            a = data_augmentation(image_batch)\n",
        "            b = data_augmentation(image_batch)\n",
        "\n",
        "            loss = train_step(a, b, model, optimizer, criterion, temperature)\n",
        "            step_wise_loss.append(loss)\n",
        "\n",
        "        epoch_wise_loss.append(np.mean(step_wise_loss))\n",
        "        wandb.log({\"nt_xentloss\": np.mean(step_wise_loss)})\n",
        "        \n",
        "        if epoch % 10 == 0:\n",
        "            print(\"epoch: {} loss: {:.3f}\".format(epoch + 1, np.mean(step_wise_loss)))\n",
        "\n",
        "    return epoch_wise_loss, model"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "atu93xMCgBYj",
        "colab_type": "text"
      },
      "source": [
        "## Training"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab_type": "code",
        "id": "QKAoeXwJZ5S1",
        "colab": {}
      },
      "source": [
        "criterion = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True, \n",
        "                                                          reduction=tf.keras.losses.Reduction.SUM)\n",
        "decay_steps = 1000\n",
        "lr_decayed_fn = tf.keras.experimental.CosineDecay(\n",
        "    initial_learning_rate=0.1, decay_steps=decay_steps)\n",
        "optimizer = tf.keras.optimizers.SGD(lr_decayed_fn)\n",
        "\n",
        "resnet_simclr_2 = get_resnet_simclr(256, 128, 50)\n",
        "\n",
        "epoch_wise_loss, resnet_simclr  = train_simclr(resnet_simclr_2, train_ds, optimizer, criterion,\n",
        "                 temperature=0.1, epochs=200)\n",
        "\n",
        "with plt.xkcd():\n",
        "    plt.plot(epoch_wise_loss)\n",
        "    plt.title(\"tau = 0.1, h1 = 256, h2 = 128, h3 = 50\")\n",
        "    plt.show()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "HXCy84vvgBYo",
        "colab_type": "text"
      },
      "source": [
        "## Store the weights to a safe place :)"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab_type": "code",
        "id": "D-zrHcxhaSUh",
        "colab": {}
      },
      "source": [
        "from google.colab import auth\n",
        "auth.authenticate_user()\n",
        "\n",
        "PROJECT_ID = \"fast-ai-exploration\" \n",
        "!gcloud config set project $PROJECT_ID"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab_type": "code",
        "id": "_OyIyB-kkHuN",
        "colab": {}
      },
      "source": [
        "import datetime\n",
        "filename = datetime.datetime.now().strftime(\"%Y%m%d-%H%M%S\") + \"resnet_simclr.h5\"\n",
        "filename"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab_type": "code",
        "id": "qNzIZ0e4kMGk",
        "colab": {}
      },
      "source": [
        "#!gsutil mb gs://resnet_simclr_imagenet\n",
        "resnet_simclr.save_weights(filename)\n",
        "!gsutil -m cp 20200508-134915resnet_simclr.h5 gs://resnet_simclr_imagenet/"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "vUrvFEpesI4Q",
        "colab_type": "text"
      },
      "source": [
        "## In case saving to a GCS bucket is not possible"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "CdkBEeB5sMU4",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "resnet_simclr.save_weights(filename)\n",
        "wandb.save(filename)"
      ],
      "execution_count": null,
      "outputs": []
    }
  ]
}