regular-classification.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "name": "regular-classification.ipynb",
      "provenance": [],
      "collapsed_sections": [],
      "machine_shape": "hm"
    },
    "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.7.12"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "HJ0S9kEasyjc"
      },
      "source": [
        "## Imports and setup"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "a60ZN1XJqjqs"
      },
      "source": [
        "!pip install -q tensorflow-addons"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "KclKS2uSqsTn"
      },
      "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)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "2CAdE6uZs1bg"
      },
      "source": [
        "## Hyperparameters"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "nGNq1YE9quPL"
      },
      "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",
        "\n",
        "# ENCODER and DECODER\n",
        "LAYER_NORM_EPS = 1e-6\n",
        "ENC_PROJECTION_DIM = 128\n",
        "ENC_NUM_HEADS = 4\n",
        "ENC_LAYERS = 3\n",
        "ENC_TRANSFORMER_UNITS = [\n",
        "    ENC_PROJECTION_DIM * 2,\n",
        "    ENC_PROJECTION_DIM,\n",
        "]  # Size of the transformer layers."
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "82XzhKTus3Ol"
      },
      "source": [
        "## CIFAR-10 dataset loading and preparation"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "fNENQQyJqwDY"
      },
      "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)}\")"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "UkKfe5hWq19S"
      },
      "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",
        "        [\n",
        "            layers.Rescaling(1 / 255.0),\n",
        "            layers.Resizing(IMAGE_SIZE, IMAGE_SIZE),\n",
        "        ],\n",
        "        name=\"test_data_augmentation\",\n",
        "    )\n",
        "    return model"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "LlCeqaYaqxWu"
      },
      "source": [
        "def prepare_data(images, labels, is_train=True):\n",
        "    if is_train:\n",
        "        augmentation_model = get_train_augmentation_model()\n",
        "    else:\n",
        "        augmentation_model = get_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)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "cXWdbsOjs6ek"
      },
      "source": [
        "## Patchify layers"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "H5rZ-br7q3lL"
      },
      "source": [
        "class Patches(layers.Layer):\n",
        "    def __init__(self, patch_size=PATCH_SIZE):\n",
        "        super(Patches, self).__init__()\n",
        "        self.patch_size = patch_size\n",
        "\n",
        "    def call(self, images):\n",
        "        batch_size = tf.shape(images)[0]\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",
        "        patch_dims = patches.shape[-1]\n",
        "        patches = tf.reshape(patches, [batch_size, -1, patch_dims])\n",
        "        return patches"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "hJ40Eh_9q9F-"
      },
      "source": [
        "class PatchEncoder(layers.Layer):\n",
        "    def __init__(self, num_patches=NUM_PATCHES, projection_dim=ENC_PROJECTION_DIM):\n",
        "        super(PatchEncoder, self).__init__()\n",
        "        self.num_patches = num_patches\n",
        "        self.projection = layers.Dense(units=projection_dim)\n",
        "        self.position_embedding = layers.Embedding(\n",
        "            input_dim=num_patches, output_dim=projection_dim\n",
        "        )\n",
        "\n",
        "    def call(self, patch):\n",
        "        positions = tf.range(start=0, limit=self.num_patches, delta=1)\n",
        "        encoded = self.projection(patch) + self.position_embedding(positions)\n",
        "        return encoded"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "HSgz58Uys_wN"
      },
      "source": [
        "## ViT model utility"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "VGRHNeUW7wWN"
      },
      "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"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "WDtU--GgrPGG"
      },
      "source": [
        "def create_vit_classifier():\n",
        "    inputs = layers.Input(shape=(IMAGE_SIZE, IMAGE_SIZE, 3))\n",
        "    # Create patches.\n",
        "    patches = Patches()(inputs)\n",
        "    # Encode patches.\n",
        "    encoded_patches = PatchEncoder()(patches)\n",
        "\n",
        "    # Create multiple layers of the Transformer block.\n",
        "    for _ in range(ENC_LAYERS):\n",
        "        # Layer normalization 1.\n",
        "        x1 = layers.LayerNormalization(epsilon=LAYER_NORM_EPS)(encoded_patches)\n",
        "        # Create a multi-head attention layer.\n",
        "        attention_output = layers.MultiHeadAttention(\n",
        "            num_heads=ENC_NUM_HEADS, key_dim=ENC_PROJECTION_DIM, dropout=0.1\n",
        "        )(x1, x1)\n",
        "        # Skip connection 1.\n",
        "        x2 = layers.Add()([attention_output, encoded_patches])\n",
        "        # Layer normalization 2.\n",
        "        x3 = layers.LayerNormalization(epsilon=LAYER_NORM_EPS)(x2)\n",
        "        # MLP.\n",
        "        x3 = mlp(x3, hidden_units=ENC_TRANSFORMER_UNITS, dropout_rate=0.1)\n",
        "        # Skip connection 2.\n",
        "        encoded_patches = layers.Add()([x3, x2])\n",
        "\n",
        "    # Create a [batch_size, projection_dim] tensor.\n",
        "    representation = layers.LayerNormalization(epsilon=LAYER_NORM_EPS)(encoded_patches)\n",
        "    representation = layers.GlobalAveragePooling1D()(representation)\n",
        "    # Classify outputs.\n",
        "    outputs = layers.Dense(NUM_CLASSES, activation=\"softmax\")(representation)\n",
        "    # Create the Keras model.\n",
        "    model = keras.Model(inputs=inputs, outputs=outputs)\n",
        "    return model"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ihd604hFtCcg"
      },
      "source": [
        "## LR scheduler"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "v2wkvqhlsQ-M"
      },
      "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",
        "        )"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "5PtvlXWYsXbD"
      },
      "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",
        ")"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "bnQzF-v3tE0X"
      },
      "source": [
        "## Train and evaluate model"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "HT-gf8mpsZ5F"
      },
      "source": [
        "optimizer = tfa.optimizers.AdamW(\n",
        "    learning_rate=scheduled_lrs,\n",
        "    weight_decay=WEIGHT_DECAY\n",
        ")\n",
        "\n",
        "vit_model = create_vit_classifier()\n",
        "vit_model.compile(\n",
        "    optimizer=optimizer,\n",
        "    loss=\"sparse_categorical_crossentropy\",\n",
        "    metrics=[\"accuracy\"]\n",
        ")\n",
        "vit_model.fit(train_ds, validation_data=val_ds, epochs=EPOCHS)\n",
        "\n",
        "loss, accuracy = vit_model.evaluate(test_ds)\n",
        "accuracy = round(accuracy * 100, 2)\n",
        "print(f\"Accuracy on the test set: {accuracy}%.\")"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "lfbsfkGv7wWP"
      },
      "source": [
        "vit_model.save(f\"classification_vit_model@acc_{accuracy}\", include_optimizer=False)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "0ZcA2gkN7wWQ"
      },
      "source": [
        "## References\n",
        "\n",
        "* https://keras.io/examples/vision/image_classification_with_vision_transformer/"
      ]
    }
  ]
}