ViT-demo.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import tensorflow as tf\n",
    "import tensorflow_addons as tfa\n",
    "import numpy as np\n",
    "import utils\n",
    "import matplotlib.pyplot as plt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# hyperparameters taken from\n",
    "# https://github.com/keras-team/keras-io/blob/master/examples/vision/image_classification_with_vision_transformer.py#L54\n",
    "learning_rate = 0.001\n",
    "weight_decay = 0.0001\n",
    "batch_size = 256\n",
    "num_epochs = 100\n",
    "image_size = 72  # We'll resize input images to this size\n",
    "patch_size = 6  # Size of the patches to be extract from the input images\n",
    "projection_dim = 64\n",
    "num_heads = 4\n",
    "transformer_units = [\n",
    "    projection_dim * 2,\n",
    "    projection_dim,\n",
    "]\n",
    "transformer_layers = 8\n",
    "mlp_head_units = [2048, 1024]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_dict = utils.unpickle(\"./cifar-100-python/train\")\n",
    "test_dict = utils.unpickle(\"./cifar-100-python/test\")\n",
    "# Keys are unaccessible because they are in bytestring format\n",
    "def decode_keys(data):\n",
    "    decoded = []\n",
    "    for i in data.keys():\n",
    "        decoded.append(i.decode(\"utf-8\"))\n",
    "    return dict(zip(decoded, data.values()))\n",
    "train_dict_decoded = decode_keys(train_dict)\n",
    "test_dict_decoded = decode_keys(test_dict)\n",
    "print(train_dict_decoded.keys())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Explore data format\n",
    "print(train_dict_decoded[\"data\"][0].shape)\n",
    "print(train_dict_decoded[\"fine_labels\"][0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create a function to isolate each color channel and build the image\n",
    "def prepare_dict(data):\n",
    "    from copy import copy\n",
    "    data_copy = copy(data)\n",
    "    processed_imgs = []\n",
    "    for i in data_copy[\"data\"]: #bruh\n",
    "        processed_imgs.append(np.reshape(i, (3, 32, 32)).transpose(1, 2, 0))\n",
    "    data_copy[\"data\"] = processed_imgs\n",
    "    data_copy[\"label\"] = data_copy[\"fine_labels\"]\n",
    "    del data_copy[\"filenames\"]\n",
    "    del data_copy[\"batch_label\"]\n",
    "    del data_copy[\"coarse_labels\"]\n",
    "    del data_copy[\"fine_labels\"]\n",
    "    return data_copy\n",
    "\n",
    "train_data = prepare_dict(train_dict_decoded)\n",
    "test_data = prepare_dict(test_dict_decoded)\n",
    "print(train_data[\"data\"][0].shape)\n",
    "print(train_data[\"label\"][0])\n",
    "print(train_data[\"data\"][0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Creating train and test dataset\n",
    "X_train, y_train = np.array(train_data[\"data\"]), train_data[\"label\"]\n",
    "X_test, y_test = np.array(test_data[\"data\"]), test_data[\"label\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Loading metadata\n",
    "meta_dict = decode_keys(utils.unpickle(\"./cifar-100-python/meta\"))\n",
    "class_names = [i.decode(\"utf-8\") for i in meta_dict[\"fine_label_names\"]]\n",
    "print(class_names[:35])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Each index corresponds to the \"label\" value in the dataset\n",
    "# Visualize 9 images\n",
    "def visualize_data(train_images, train_labels):\n",
    "    import random\n",
    "    random_idx = random.sample(list(range(len(train_labels))), 9)\n",
    "    plt.figure(figsize=(9, 9))\n",
    "    for i in range(9):\n",
    "        subp = int(\"33\" + str(i + 1))\n",
    "        plt.subplot(subp)\n",
    "        class_ = train_labels[random_idx[i]]\n",
    "        title = \"Class: \" + str(class_) + \", Label:\" + class_names[class_]\n",
    "        plt.title(title)\n",
    "        plt.axis(\"off\")\n",
    "        plt.grid(False)\n",
    "        plt.imshow(train_images[random_idx[i]])\n",
    "visualize_data(X_train, y_train)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# preprocess layers\n",
    "norm = tf.keras.layers.Normalization()\n",
    "res = tf.keras.layers.Resizing(image_size, image_size)\n",
    "X_train = norm(res(X_train))\n",
    "X_test = norm(res(X_test))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# create efficient data pipelines\n",
    "train_ds = tf.data.Dataset.from_tensor_slices((X_train, y_train)).shuffle(10000000).batch(batch_size).prefetch(tf.data.AUTOTUNE)\n",
    "test_ds = tf.data.Dataset.from_tensor_slices((X_test, y_test)).batch(batch_size).prefetch(tf.data.AUTOTUNE)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from model import ViT\n",
    "# create the model\n",
    "model = ViT(image_size=image_size,\n",
    "    patch_size=patch_size,\n",
    "    num_classes=len(meta_dict[\"fine_label_names\"]),\n",
    "    dim=projection_dim,\n",
    "    depth=transformer_layers,\n",
    "    attention_heads=num_heads,\n",
    "    transformer_units=transformer_units,\n",
    "    mlp_head_units=mlp_head_units\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# function to run an experiment\n",
    "def run_experiment(model):\n",
    "    optimizer = tfa.optimizers.AdamW(\n",
    "        learning_rate=learning_rate, weight_decay=weight_decay\n",
    "    )\n",
    "\n",
    "    model.compile(\n",
    "        optimizer=optimizer,\n",
    "        loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),\n",
    "        metrics=[\n",
    "            tf.keras.metrics.SparseCategoricalAccuracy(name=\"accuracy\"),\n",
    "            tf.keras.metrics.SparseTopKCategoricalAccuracy(5, name=\"top-5-accuracy\"),\n",
    "        ],\n",
    "    )\n",
    "\n",
    "    checkpoint_filepath = \"/tmp/checkpoint\"\n",
    "    checkpoint_callback = tf.keras.callbacks.ModelCheckpoint(\n",
    "        checkpoint_filepath,\n",
    "        monitor=\"val_accuracy\",\n",
    "        save_best_only=True,\n",
    "        save_weights_only=True,\n",
    "    )\n",
    "\n",
    "    history = model.fit(\n",
    "        train_ds,\n",
    "        batch_size=batch_size,\n",
    "        epochs=num_epochs,\n",
    "        validation_data=test_ds,\n",
    "        callbacks=[checkpoint_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)}%\")\n",
    "\n",
    "    return history\n",
    "model.build((None, 32, 32, 3))\n",
    "print(model.summary())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Call the ViT\n",
    "run_experiment(model=model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.10.6 ('tensorflow')",
   "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.10.8"
  },
  "orig_nbformat": 4,
  "vscode": {
   "interpreter": {
    "hash": "4bd624a0593993fe43ac4046b27b898fb2ef75c21c08f81e89e64ea0f51df676"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}