notebooks/01_Model_Selection_vit-base-patch16-224.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "643e3192",
   "metadata": {},
   "source": [
    "# Tensorflow Image Classifier\n",
    "\n",
    "## Vision Transformer\n",
    "\n",
    "> [google/vit-base-patch16-224](https://arxiv.org/pdf/2010.11929.pdf)\n",
    "\n",
    "\n",
    "_\"The Vision Transformer (ViT) is a transformer encoder model (BERT-like) pretrained on a large collection of images in a supervised fashion, namely ImageNet-21k, at a resolution of 224x224 pixels. Next, the model was fine-tuned on ImageNet (also referred to as ILSVRC2012), a dataset comprising 1 million images and 1,000 classes, also at resolution 224x224.\"_\n",
    "\n",
    "\n",
    "This is using the [same transformer](https://huggingface.co/docs/transformers/model_doc/vit) as before - but is now pre-trained on a large dataset. Let's see if this can beat the results I received from the not-pre-trained ViT before:\n",
    "\n",
    "* `loss: 0.9803 - accuracy: 0.5443 - topk_accuracy: 0.8090` (self-build: num_attention_heads=4,num_hidden_layers=2, hidden_size=144)\n",
    "* `loss: 0.9404 - accuracy: 0.5610 - topk_accuracy: 0.8385` (hf model: num_attention_heads=4,num_hidden_layers=2, hidden_size=144)\n",
    "* `loss: 0.9572 - accuracy: 0.5399 - topk_accuracy: 0.8306` (hf model: num_attention_heads=8,num_hidden_layers=4, hidden_size=144)\n",
    "* `loss: 1.0289 - accuracy: 0.5000 - topk_accuracy: 0.7761` (hf model: num_attention_heads=8,num_hidden_layers=4, hidden_size=768)\n",
    "    * The accuracy here is lower - but I massively reduced the learning rate and the accuracy was still rising steadily \n",
    "* `loss: 0.5340 - accuracy: 0.7730 - topk_accuracy: 0.9315` (hf model with pre-trained weights: num_attention_heads=8,num_hidden_layers=4, hidden_size=768)\n",
    "* `loss: 0.4596 - accuracy: 0.8068 - topk_accuracy: 0.9495` (re-run as above with 40 instead of 20 epochs)\n",
    "\n",
    "For the training with pretrained weights I increased the learning-rate again by a factor of 10. But I had to down-configure the model to prevent out-of-memory errors (configuration as seen below). So there is still room to improve the results on better hardware and with some tinkering. Still - this pre-trained model performs significantly better than the untrained transformer and already comes close to rival the CNN solution (MobileNetV3Small) used before:\n",
    "\n",
    "* `loss: 0.3906 - accuracy: 0.8455 - topk_accuracy: 0.9627` - __MobileNetV3Small__ as reference"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0ae495d4",
   "metadata": {},
   "outputs": [],
   "source": [
    "@misc{wu2020visual,\n",
    "      title={Visual Transformers: Token-based Image Representation and Processing for Computer Vision}, \n",
    "      author={Bichen Wu and Chenfeng Xu and Xiaoliang Dai and Alvin Wan and Peizhao Zhang and Zhicheng Yan and Masayoshi Tomizuka and Joseph Gonzalez and Kurt Keutzer and Peter Vajda},\n",
    "      year={2020},\n",
    "      eprint={2006.03677},\n",
    "      archivePrefix={arXiv},\n",
    "      primaryClass={cs.CV}\n",
    "}\n",
    "\n",
    "@misc{wu2020visual,\n",
    "      title={Visual Transformers: Token-based Image Representation and Processing for Computer Vision}, \n",
    "      author={Bichen Wu and Chenfeng Xu and Xiaoliang Dai and Alvin Wan and Peizhao Zhang and Zhicheng Yan and Masayoshi Tomizuka and Joseph Gonzalez and Kurt Keutzer and Peter Vajda},\n",
    "      year={2020},\n",
    "      eprint={2006.03677},\n",
    "      archivePrefix={arXiv},\n",
    "      primaryClass={cs.CV}\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "db774445",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(tf.test.is_gpu_available)\n",
    "print(tf.__version__)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fa5cc00f-82d7-43f1-8635-dcb0dc417bf3",
   "metadata": {},
   "outputs": [],
   "source": [
    "!apt-get install libgl1 -y"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "54818db4",
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install opencv-python scikit-learn seaborn transformers"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4b11c094",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "import cv2 as cv\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "from sklearn.metrics import (\n",
    "    classification_report,\n",
    "    confusion_matrix,\n",
    "    ConfusionMatrixDisplay)\n",
    "import tensorflow as tf\n",
    "from tensorflow.keras import Sequential\n",
    "from tensorflow.keras.layers import (\n",
    "    Conv2D,\n",
    "    BatchNormalization,\n",
    "    LayerNormalization,\n",
    "    Dense,\n",
    "    Input,\n",
    "    Embedding,\n",
    "    MultiHeadAttention,\n",
    "    Layer,\n",
    "    Add,\n",
    "    Resizing,\n",
    "    Rescaling,\n",
    "    Permute,\n",
    "    Flatten,\n",
    "    RandomFlip,\n",
    "    RandomRotation,\n",
    "    RandomContrast,\n",
    "    RandomBrightness\n",
    ")\n",
    "from tensorflow.keras.losses import CategoricalCrossentropy\n",
    "from tensorflow.keras.metrics import CategoricalAccuracy, TopKCategoricalAccuracy\n",
    "from tensorflow.keras.models import Model\n",
    "from tensorflow.keras.optimizers import Adam\n",
    "from tensorflow.keras.utils import image_dataset_from_directory\n",
    "\n",
    "from transformers import ViTConfig, ViTModel, AutoImageProcessor, TFViTModel"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "365c32e1",
   "metadata": {},
   "outputs": [],
   "source": [
    "SEED = 42\n",
    "LABELS = ['Gladiolus', 'Adenium', 'Alpinia_Purpurata', 'Alstroemeria', 'Amaryllis', 'Anthurium_Andraeanum', 'Antirrhinum', 'Aquilegia', 'Billbergia_Pyramidalis', 'Cattleya', 'Cirsium', 'Coccinia_Grandis', 'Crocus', 'Cyclamen', 'Dahlia', 'Datura_Metel', 'Dianthus_Barbatus', 'Digitalis', 'Echinacea_Purpurea', 'Echinops_Bannaticus', 'Fritillaria_Meleagris', 'Gaura', 'Gazania', 'Gerbera', 'Guzmania', 'Helianthus_Annuus', 'Iris_Pseudacorus', 'Leucanthemum', 'Malvaceae', 'Narcissus_Pseudonarcissus', 'Nerine', 'Nymphaea_Tetragona', 'Paphiopedilum', 'Passiflora', 'Pelargonium', 'Petunia', 'Platycodon_Grandiflorus', 'Plumeria', 'Poinsettia', 'Primula', 'Protea_Cynaroides', 'Rose', 'Rudbeckia', 'Strelitzia_Reginae', 'Tropaeolum_Majus', 'Tussilago', 'Viola', 'Zantedeschia_Aethiopica']\n",
    "NLABELS = len(LABELS)\n",
    "BATCH_SIZE = 32\n",
    "SIZE = 224 # 256\n",
    "EPOCHS = 40\n",
    "LR = 5e-6 # default 0.001\n",
    "HIDDEN_SIZE = 768 # default 768\n",
    "NHEADS = 8 # default 12\n",
    "NLAYERS = 4 # default 12"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4ef778cd",
   "metadata": {},
   "source": [
    "### Dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0f753146",
   "metadata": {},
   "outputs": [],
   "source": [
    "train_directory = '../dataset/Flower_Dataset/split/train'\n",
    "test_directory = '../dataset/Flower_Dataset/split/val'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0ed88828",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "train_dataset = image_dataset_from_directory(\n",
    "    train_directory,\n",
    "    labels='inferred',\n",
    "    label_mode='categorical',\n",
    "    class_names=LABELS,\n",
    "    color_mode='rgb',\n",
    "    batch_size=BATCH_SIZE,\n",
    "    image_size=(SIZE, SIZE),\n",
    "    shuffle=True,\n",
    "    seed=SEED,\n",
    "    interpolation='bilinear',\n",
    "    follow_links=False,\n",
    "    crop_to_aspect_ratio=False\n",
    ")\n",
    "\n",
    "# Found 9206 files belonging to 48 classes.\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "02dda8be",
   "metadata": {},
   "outputs": [],
   "source": [
    "test_dataset = image_dataset_from_directory(\n",
    "    test_directory,\n",
    "    labels='inferred',\n",
    "    label_mode='categorical',\n",
    "    class_names=LABELS,\n",
    "    color_mode='rgb',\n",
    "    batch_size=BATCH_SIZE,\n",
    "    image_size=(SIZE, SIZE),\n",
    "    shuffle=True,\n",
    "    seed=SEED\n",
    ")\n",
    "\n",
    "# Found 3090 files belonging to 48 classes."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8b670460",
   "metadata": {},
   "outputs": [],
   "source": [
    "data_augmentation = Sequential([\n",
    "    RandomRotation(factor = (0.25),),\n",
    "    RandomFlip(mode='horizontal',),\n",
    "    RandomContrast(factor=0.1),\n",
    "    RandomBrightness(0.1)\n",
    "])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a7921c41-8ff6-4485-a71f-5c7ed6c92a62",
   "metadata": {},
   "outputs": [],
   "source": [
    "resize_rescale_reshape = Sequential([\n",
    "    Resizing(SIZE, SIZE),\n",
    "    Rescaling(1./255),\n",
    "    # transformer expects image shape (3,224,224)\n",
    "    Permute((3,1,2))\n",
    "])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3bcc45ae",
   "metadata": {},
   "outputs": [],
   "source": [
    "training_dataset = (\n",
    "    train_dataset\n",
    "    .map(lambda image, label: (data_augmentation(image), label))\n",
    "    .prefetch(tf.data.AUTOTUNE)\n",
    ")\n",
    "\n",
    "\n",
    "testing_dataset = (\n",
    "    test_dataset.prefetch(\n",
    "        tf.data.AUTOTUNE\n",
    "    )\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d7f3770b",
   "metadata": {},
   "source": [
    "### ViT Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "56ec31f1",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Initializing a ViT vit-base-patch16-224 style configuration\n",
    "configuration = ViTConfig(\n",
    "    image_size=SIZE,\n",
    "    hidden_size=HIDDEN_SIZE,\n",
    "    num_attention_heads=NHEADS,\n",
    "    num_hidden_layers=NLAYERS\n",
    ")\n",
    "\n",
    "# Initializing a model with random weights from the vit-base-patch16-224 style configuration\n",
    "# base_model = TFViTModel(configuration)\n",
    "\n",
    "# use pretrained weights for the model instead\n",
    "base_model = TFViTModel.from_pretrained(\"google/vit-base-patch16-224-in21k\", config=configuration)\n",
    "\n",
    "# Accessing the model configuration\n",
    "configuration = base_model.config"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "18e44112-7d14-41db-9c0a-15162f2a5dd1",
   "metadata": {},
   "outputs": [],
   "source": [
    "configuration"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d94afcd8-ca1a-490d-9e22-9e1c6d43f4f6",
   "metadata": {},
   "outputs": [],
   "source": [
    "inputs = Input(shape=(224,224,3))\n",
    "x = resize_rescale_reshape(inputs)\n",
    "x = base_model.vit(x)[0][:,0,:]\n",
    "output = Dense(NLABELS, activation='softmax')(x)\n",
    "\n",
    "vit_model = Model(inputs=inputs, outputs=output)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3fbf81dd-f624-4ba5-b254-17b43c87a20b",
   "metadata": {},
   "outputs": [],
   "source": [
    "vit_model.summary()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c5fc4f3f-a440-40e3-8e11-a2228df28d82",
   "metadata": {},
   "outputs": [],
   "source": [
    "# testing the pretrained model\n",
    "test_image = cv.imread('../dataset/snapshots/Viola_Tricolor.jpg')\n",
    "test_image = cv.resize(test_image, (SIZE, SIZE))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6536e3b1-56d4-4ef3-8b25-a1555f06ee19",
   "metadata": {},
   "outputs": [],
   "source": [
    "vit_model(tf.expand_dims(test_image, axis = 0))\n",
    "# numpy=array([[0.01901147, 0.02026679, 0.02391427, 0.01970932, 0.01464635,\n",
    "        # 0.0257492 , 0.01927904, 0.01793713, 0.01944521, 0.01977614,\n",
    "        # 0.02793077, 0.02291007, 0.02077055, 0.02195414, 0.01900317,\n",
    "        # 0.01640951, 0.0187414 , 0.02054461, 0.01795707, 0.01564359,\n",
    "        # 0.02500662, 0.02195591, 0.02427697, 0.01805321, 0.01870451,\n",
    "        # 0.01892176, 0.01930878, 0.02687679, 0.02315602, 0.02085607,\n",
    "        # 0.01970802, 0.02608317, 0.02246164, 0.01824699, 0.02068511,\n",
    "        # 0.0230596 , 0.02106061, 0.02080243, 0.02133719, 0.02659844,\n",
    "        # 0.02275858, 0.02423375, 0.01562007, 0.01791171, 0.02137934,\n",
    "        # 0.02457437, 0.01662222, 0.01814036]], dtype=float32)>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9ec69a58",
   "metadata": {},
   "source": [
    "### Model Training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cc980d15",
   "metadata": {},
   "outputs": [],
   "source": [
    "loss_function = CategoricalCrossentropy()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "79a30969",
   "metadata": {},
   "outputs": [],
   "source": [
    "metrics = [CategoricalAccuracy(name='accuracy')]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2c4c9de5",
   "metadata": {},
   "outputs": [],
   "source": [
    "vit_model.compile(\n",
    "    optimizer = Adam(learning_rate = LR),\n",
    "    loss = loss_function,\n",
    "    metrics = metrics\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5786148c",
   "metadata": {},
   "outputs": [],
   "source": [
    "vit_history = vit_model.fit(\n",
    "    training_dataset,\n",
    "    validation_data = testing_dataset,\n",
    "    epochs = EPOCHS,\n",
    "    verbose = 1\n",
    ")\n",
    "\n",
    "# loss: 0.4731\n",
    "# accuracy: 0.9034\n",
    "# val_loss: 0.5814\n",
    "# val_accuracy: 0.8663"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c682a3a0",
   "metadata": {},
   "outputs": [],
   "source": [
    "vit_model.evaluate(testing_dataset)\n",
    "# loss: 0.5814 - accuracy: 0.8663"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "58bffe28",
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.plot(vit_history.history['loss'])\n",
    "plt.plot(vit_history.history['val_loss'])\n",
    "plt.title('Model Loss')\n",
    "plt.ylabel('loss')\n",
    "plt.xlabel('epoch')\n",
    "plt.legend(['train_loss', 'val_loss'])\n",
    "\n",
    "plt.savefig('assets/ViT_01.webp', bbox_inches='tight')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bb0419a3",
   "metadata": {},
   "source": [
    "![Building a Tensorflow VIT](./assets/ViT_01.webp)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ddc6089b",
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.plot(vit_history.history['accuracy'])\n",
    "plt.plot(vit_history.history['val_accuracy'])\n",
    "plt.title('Model Accuracy')\n",
    "plt.ylabel('accuracy')\n",
    "plt.xlabel('epoch')\n",
    "plt.legend(['train_accuracy', 'val_accuracy'])\n",
    "\n",
    "plt.savefig('assets/ViT_02.webp', bbox_inches='tight')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "06fc766d",
   "metadata": {},
   "source": [
    "![Building a Tensorflow VIT](./assets/ViT_02.webp)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "091823e3-9ae1-4c67-b94a-60e5e49306ba",
   "metadata": {},
   "outputs": [],
   "source": [
    "test_image_bgr = cv.imread('../dataset/snapshots/Viola_Tricolor.jpg')\n",
    "test_image_resized = cv.resize(test_image_bgr, (SIZE, SIZE))\n",
    "test_image_rgb = cv.cvtColor(test_image_resized, cv.COLOR_BGR2RGB)\n",
    "img = tf.constant(test_image_rgb, dtype=tf.float32)\n",
    "img = tf.expand_dims(img, axis=0)\n",
    "\n",
    "probs = vit_model(img).numpy()\n",
    "label = LABELS[tf.argmax(probs, axis=1).numpy()[0]]\n",
    "\n",
    "print(label, str(probs[0]))\n",
    "\n",
    "plt.imshow(test_image_rgb)\n",
    "plt.title(label)\n",
    "plt.axis('off')\n",
    "        \n",
    "plt.savefig('assets/ViT_Prediction_01.webp', bbox_inches='tight')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4f548907",
   "metadata": {},
   "source": [
    "![Building a Tensorflow VIT](./assets/ViT_Prediction_01.webp)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8090c898-8cd3-4b2f-b0be-3ef0a0b05b59",
   "metadata": {},
   "outputs": [],
   "source": [
    "test_image_bgr = cv.imread('../dataset/snapshots/Strelitzia.jpg')\n",
    "test_image_resized = cv.resize(test_image_bgr, (SIZE, SIZE))\n",
    "test_image_rgb = cv.cvtColor(test_image_resized, cv.COLOR_BGR2RGB)\n",
    "img = tf.constant(test_image_rgb, dtype=tf.float32)\n",
    "img = tf.expand_dims(img, axis=0)\n",
    "\n",
    "probs = vit_model(img).numpy()\n",
    "label = LABELS[tf.argmax(probs, axis=1).numpy()[0]]\n",
    "\n",
    "print(label, str(probs[0]))\n",
    "\n",
    "plt.imshow(test_image_rgb)\n",
    "plt.title(label)\n",
    "plt.axis('off')\n",
    "        \n",
    "plt.savefig('assets/ViT_Prediction_02.webp', bbox_inches='tight')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9d6b1957",
   "metadata": {},
   "source": [
    "![Building a Tensorflow VIT](./assets/ViT_Prediction_02.webp)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ab7be81c-c495-432a-ad4d-54d3be24ed40",
   "metadata": {},
   "outputs": [],
   "source": [
    "test_image_bgr = cv.imread('../dataset/snapshots/Water_Lilly.jpg')\n",
    "test_image_resized = cv.resize(test_image_bgr, (SIZE, SIZE))\n",
    "test_image_rgb = cv.cvtColor(test_image_resized, cv.COLOR_BGR2RGB)\n",
    "img = tf.constant(test_image_rgb, dtype=tf.float32)\n",
    "img = tf.expand_dims(img, axis=0)\n",
    "\n",
    "probs = vit_model(img).numpy()\n",
    "label = LABELS[tf.argmax(probs, axis=1).numpy()[0]]\n",
    "\n",
    "print(label, str(probs[0]))\n",
    "\n",
    "plt.imshow(test_image_rgb)\n",
    "plt.title(label)\n",
    "plt.axis('off')\n",
    "        \n",
    "plt.savefig('assets/ViT_Prediction_03.webp', bbox_inches='tight')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "09228b3b",
   "metadata": {},
   "source": [
    "![Building a Tensorflow VIT](./assets/ViT_Prediction_03.webp)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4fcd57ac",
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.figure(figsize=(16,16))\n",
    "\n",
    "for images, labels in testing_dataset.take(1):\n",
    "    for i in range(16):\n",
    "        ax = plt.subplot(4,4,i+1)\n",
    "        true = \"True: \" + LABELS[tf.argmax(labels[i], axis=0).numpy()]\n",
    "        pred = \"Predicted: \" + LABELS[\n",
    "            tf.argmax(vit_model(tf.expand_dims(images[i], axis=0)).numpy(), axis=1).numpy()[0]\n",
    "        ]\n",
    "        plt.title(\n",
    "           true  + \"\\n\" + pred\n",
    "        )\n",
    "        plt.imshow(images[i]/255.)\n",
    "        plt.axis('off')\n",
    "        \n",
    "plt.savefig('assets/ViT_03.webp', bbox_inches='tight')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "de9856ee",
   "metadata": {},
   "source": [
    "![Building a Tensorflow VIT](./assets/ViT_03.webp)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f70875ab",
   "metadata": {},
   "outputs": [],
   "source": [
    "y_pred = []\n",
    "y_test = []\n",
    "\n",
    "for img, label in testing_dataset:\n",
    "    y_pred.append(vit_model(img))\n",
    "    y_test.append(label.numpy())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "527da3d7",
   "metadata": {},
   "outputs": [],
   "source": [
    "conf_mtx = ConfusionMatrixDisplay(\n",
    "    confusion_matrix=confusion_matrix(\n",
    "        np.argmax(y_test[:-1], axis=-1).flatten(),\n",
    "        np.argmax(y_pred[:-1], axis=-1).flatten()\n",
    "    ),\n",
    "    display_labels=LABELS\n",
    ")\n",
    "\n",
    "fig, ax = plt.subplots(figsize=(16,12))\n",
    "conf_mtx.plot(ax=ax, cmap='plasma', include_values=True, xticks_rotation='vertical',)\n",
    "\n",
    "plt.savefig('assets/ViT_04.webp', bbox_inches='tight')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b7ab425e",
   "metadata": {},
   "source": [
    "![Building a Tensorflow VIT](./assets/ViT_04.webp)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5de6b929-c1b7-44de-a68d-e946d2b0c9eb",
   "metadata": {},
   "source": [
    "### Saving the Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6eec4393-404a-4f4d-bb4f-983f937b258b",
   "metadata": {},
   "outputs": [],
   "source": [
    "tf.keras.saving.save_model(\n",
    "    vit_model, '../saved_model/vit_model/1', overwrite=True, save_format='tf'\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8b225eb2-1dc3-490b-ae8a-0399a4a9a897",
   "metadata": {},
   "outputs": [],
   "source": [
    "# restore the model\n",
    "restored_model = tf.keras.saving.load_model('../saved_model/vit_model/1')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5c84954e-abb5-4178-b92d-1f3c569f771c",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Check its architecture\n",
    "restored_model.summary()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "988cf5ad-8588-47e5-8d11-3158307c55b1",
   "metadata": {},
   "outputs": [],
   "source": [
    "restored_model.evaluate(testing_dataset)\n",
    "# loss: 0.5979 - accuracy: 0.8589"
   ]
  }
 ],
 "metadata": {
  "jupytext": {
   "formats": "ipynb,md"
  },
  "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.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}