ViT_MNISTdigit_v1a.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": []
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "source": [
        "Loading the datasets"
      ],
      "metadata": {
        "id": "PRpKnxpDjYEP"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "#pip install -U tensorflow-addons\n",
        "https://keras.io/examples/vision/image_classification_with_vision_transformer/"
      ],
      "metadata": {
        "id": "KQiyEAEXlLpA"
      }
    },
    {
      "cell_type": "code",
      "execution_count": 49,
      "metadata": {
        "id": "eUMxKT4pjTJ5"
      },
      "outputs": [],
      "source": [
        "import numpy as np\n",
        "import tensorflow as tf\n",
        "from tensorflow import keras\n",
        "from tensorflow.keras import layers\n",
        "import tensorflow_addons as tfa\n",
        "import matplotlib.pyplot as plt"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# Define a function to change the MNIST label dimension\n",
        "def change_to_right(wrong_labels):\n",
        "    right_labels=[]\n",
        "    for x in wrong_labels:\n",
        "        for i in range(0,len(wrong_labels[0])):\n",
        "            if x[i]==1:\n",
        "                right_labels.append(i)\n",
        "    return right_labels\n",
        "\n",
        "# Model / data parameters\n",
        "num_classes = 10\n",
        "input_shape = (28, 28, 1)\n",
        "\n",
        "# Load the data and split it between train and test sets\n",
        "(x_train, y_train), (x_test, y_test) = keras.datasets.mnist.load_data()\n",
        "\n",
        "# Scale images to the [0, 1] range\n",
        "x_train = x_train.astype(\"float32\") / 255\n",
        "x_test = x_test.astype(\"float32\") / 255\n",
        "\n",
        "# Make sure images have shape (28, 28, 1)\n",
        "x_train = np.expand_dims(x_train, -1)\n",
        "x_test = np.expand_dims(x_test, -1)\n",
        "\n",
        "# Convert class vectors to right format\n",
        "y_train = tf.convert_to_tensor(np.array(change_to_right(keras.utils.to_categorical(y_train, num_classes))))\n",
        "y_test = tf.convert_to_tensor(np.array(change_to_right(keras.utils.to_categorical(y_test, num_classes))))"
      ],
      "metadata": {
        "id": "byCKXcT1kUR5"
      },
      "execution_count": 50,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Configuring the hyperparameters"
      ],
      "metadata": {
        "id": "8yV2-zlHkZIv"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "learning_rate = 0.001\n",
        "weight_decay = 0.0001\n",
        "batch_size = 2000\n",
        "num_epochs = 5"
      ],
      "metadata": {
        "id": "IZrmYd2NkbwS"
      },
      "execution_count": 51,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "image_size = 28  # We'll resize input images to this size\n",
        "patch_size = 7  # Size of the patches to be extract from the input images\n",
        "num_patches = (image_size // patch_size) ** 2\n",
        "projection_dim = 96\n",
        "num_heads = 4\n",
        "transformer_units = [\n",
        "    projection_dim * 2,\n",
        "    projection_dim,\n",
        "]  # Size of the transformer layers\n",
        "transformer_layers = 16\n",
        "mlp_head_units = [2048, 1024]  # Size of the dense layers of the final classifier"
      ],
      "metadata": {
        "id": "u-HjNjwXkdeI"
      },
      "execution_count": 52,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Data augmentation"
      ],
      "metadata": {
        "id": "KHpLWQtQkiNC"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "data_augmentation = keras.Sequential(\n",
        "    [\n",
        "        layers.Normalization(),\n",
        "        #layers.Resizing(image_size, image_size),\n",
        "        #layers.RandomFlip(\"horizontal\"),\n",
        "        layers.RandomRotation(factor=0.02),\n",
        "        layers.RandomZoom(\n",
        "            height_factor=0.2, width_factor=0.2\n",
        "        ),\n",
        "    ],\n",
        "    name=\"data_augmentation\",\n",
        ")\n",
        "# Compute the mean and the variance of the training data for normalization.\n",
        "data_augmentation.layers[0].adapt(x_train)"
      ],
      "metadata": {
        "id": "rPkGLOqGkhW4"
      },
      "execution_count": 53,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "def mlp(x, hidden_units, dropout_rate):\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\n",
        "\n",
        "class Patches(layers.Layer):\n",
        "    def __init__(self, patch_size):\n",
        "        super().__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"
      ],
      "metadata": {
        "id": "BCoG31yQluLw"
      },
      "execution_count": 54,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "class PatchEncoder(layers.Layer):\n",
        "    def __init__(self, num_patches, projection_dim):\n",
        "        super().__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\n",
        "\n",
        "def create_vit_classifier():\n",
        "    inputs = layers.Input(shape=input_shape)\n",
        "    # Augment data.\n",
        "    augmented = data_augmentation(inputs)\n",
        "    # Create patches.\n",
        "    patches = Patches(patch_size)(augmented)\n",
        "    # Encode patches.\n",
        "    encoded_patches = PatchEncoder(num_patches, projection_dim)(patches)\n",
        "\n",
        "    # Create multiple layers of the Transformer block.\n",
        "    for _ in range(transformer_layers):\n",
        "        # Layer normalization 1.\n",
        "        x1 = layers.LayerNormalization(epsilon=1e-6)(encoded_patches)\n",
        "        # Create a multi-head attention layer.\n",
        "        attention_output = layers.MultiHeadAttention(\n",
        "            num_heads=num_heads, key_dim=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=1e-6)(x2)\n",
        "        # MLP.\n",
        "        x3 = mlp(x3, hidden_units=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=1e-6)(encoded_patches)\n",
        "    representation = layers.Flatten()(representation)\n",
        "    representation = layers.Dropout(0.5)(representation)\n",
        "    # Add MLP.\n",
        "    features = mlp(representation, hidden_units=mlp_head_units, dropout_rate=0.5)\n",
        "    # Classify outputs.\n",
        "    logits = layers.Dense(num_classes)(features)\n",
        "    # Create the Keras model.\n",
        "    model = keras.Model(inputs=inputs, outputs=logits)\n",
        "\n",
        "    return model"
      ],
      "metadata": {
        "id": "3cP-npJxkpUY"
      },
      "execution_count": 55,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "plt.figure(figsize=(4, 4))\n",
        "image = x_train[np.random.choice(range(x_train.shape[0]))]\n",
        "plt.imshow(image)\n",
        "plt.axis(\"off\")\n",
        "\n",
        "resized_image = tf.image.resize(\n",
        "    tf.convert_to_tensor([image]), size=(image_size, image_size)\n",
        ")\n",
        "patches = Patches(patch_size)(resized_image)\n",
        "print(f\"Image size: {image_size} X {image_size}\")\n",
        "print(f\"Patch size: {patch_size} X {patch_size}\")\n",
        "print(f\"Patches per image: {patches.shape[1]}\")\n",
        "print(f\"Elements per patch: {patches.shape[-1]}\")\n",
        "\n",
        "n = int(np.sqrt(patches.shape[1]))\n",
        "plt.figure(figsize=(4, 4))\n",
        "for i, patch in enumerate(patches[0]):\n",
        "    ax = plt.subplot(n, n, i + 1)\n",
        "    patch_img = tf.reshape(patch, (patch_size, patch_size))\n",
        "    plt.imshow(patch_img.numpy())\n",
        "    plt.axis(\"off\")"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 742
        },
        "id": "giFthlt_lxN4",
        "outputId": "31534d4b-c895-45e1-929c-5b1766ed7543"
      },
      "execution_count": 56,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Image size: 28 X 28\n",
            "Patch size: 7 X 7\n",
            "Patches per image: 16\n",
            "Elements per patch: 49\n"
          ]
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "<Figure size 400x400 with 1 Axes>"
            ],
            "image/png": "iVBORw0KGgoAAAANSUhEUgAAAUgAAAFICAYAAAAyFGczAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjcuMSwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy/bCgiHAAAACXBIWXMAAA9hAAAPYQGoP6dpAAAHVklEQVR4nO3d36sfdB3H8XN25gxl/dpWSFCO2jQohByF/SBkmN1IEIh0YaAJRhJdaGhdmVYEGRVIJUUEFRl2IdGPixUWiIouFS2skWlQVI6MSidubt/+gHjC+yscz3Z4PK5ffPnCgef53Lz5ri4Wi8UKAP9ny0Z/AYCTlUACBIEECAIJEAQSIAgkQBBIgCCQAEEgAcLW6fCiLZeu5/cAeMkcOHHHaOcFCRAEEiAIJEAQSIAgkABBIAGCQAIEgQQIAgkQBBIgCCRAEEiAIJAAQSABgkACBIEECAIJEAQSIAgkQBBIgCCQAEEgAYJAAgSBBAgCCRAEEiAIJEAQSIAgkABBIAGCQAIEgQQIAgkQBBIgCCRAEEiAIJAAQSABgkACBIEECAIJEAQSIAgkQBBIgCCQAEEgAYJAAgSBBAgCCRAEEiAIJEAQSIAgkABBIAGCQAIEgQQIAgkQBBIgCCRAEEiAIJAAQSABgkACBIEECAIJEAQSIAgkQBBIgCCQAEEgAYJAAgSBBAgCCRAEEiAIJEAQSIAgkABBIAGCQAIEgQQIAgkQBBIgCCRAEEiAIJAAQSABgkAChK0b/QWAl87azh3r8rlPX7xnvP3He07MP3jbfLv3yoPzzx3yggQIAgkQBBIgCCRAEEiAIJAAQSABgkACBIEECAIJEJwaclI5fuHbxtsnr1qMt7u/Pt8eOev08XYZT795bbx9/rXHx9t95/1xvP3B7gPj7TK++q8nxttbf/m+8fYVj2xsorwgAYJAAgSBBAgCCRAEEiAIJEAQSIAgkABBIAGCQAKE1cViMbrBumjLpev9Xdik/v6Jd463v7rulvH26j9fMt5e9poHxtsda8+Mtz/993nj7TJ+9PD54+22v5023u58eP4rgdufeHa8XRz87Xh7Mjhw4o7RzgsSIAgkQBBIgCCQAEEgAYJAAgSBBAgCCRAEEiAIJEDwq4a8KH/51Px88K6PfXG83XfHtePtOd84PN7etuuD4+3aQ4fG2xNHjoy3y9i7cnBdPncZ89+B3Ly8IAGCQAIEgQQIAgkQBBIgCCRAEEiAIJAAQSABgkACBKeGm9zWs18/3j72mZ3j7UP7vzTenn/7/Hxwzw3zE7vjL7ww3q7OrwdX5r/7x2bnBQkQBBIgCCRAEEiAIJAAQSABgkACBIEECAIJEAQSIDg1PAU994G3j7f7b7p7vP3xzjvH230HrxhvF0v8G373g8+Mt6/a+ux4e8u9F4+353z0kfF2cezoeMupxwsSIAgkQBBIgCCQAEEgAYJAAgSBBAgCCRAEEiAIJEBwariO1nbtGm8f+/zZ4+2D7//KePvyLS8bbx89emy83bPj8Hj7wFPbx9vv3XnhePvKQ4vx9tYbvzPeXn/tlePt675wz3jLqccLEiAIJEAQSIAgkABBIAGCQAIEgQQIAgkQBBIgCCRAcGq4jt7xi7+Otz/bdWC8/c3z8z/bh354zXi7+4Z7x9uVlX+Ol3uX2C5jyxlnjLdrN87PEo9tn2/Z3LwgAYJAAgSBBAgCCRAEEiAIJEAQSIAgkABBIAGCQAIEp4br6J6r9423+8+8YLzddv+h8Xb3f5c5H9x4W848c7w9dPNbx9uji4fH2zf85Mh4y+bmBQkQBBIgCCRAEEiAIJAAQSABgkACBIEECAIJEAQSIDg1XE/3PTKeLvOHOLH8N9lQq6efPt4+/q03jbd3v+uW8faSmz453u6499Q6z2T9eEECBIEECAIJEAQSIAgkQBBIgCCQAEEgAYJAAgSBBAhODXlRVve9Zbw997bfj7efe/U3x9sPX/7x8XbHr50PsjwvSIAgkABBIAGCQAIEgQQIAgkQBBIgCCRAEEiAIJAAwanhJrd62rbx9k/fPXe8/fkFXxtvv/zU/vH205d9ZLzdcv9D4y28GF6QAEEgAYJAAgSBBAgCCRAEEiAIJEAQSIAgkABBIAGCU8OTxJM3XzDeHn/jc+PtH9777fH28RfuGm8vv/668Xb77feNtysrjy6xhfXlBQkQBBIgCCRAEEiAIJAAQSABgkACBIEECAIJEAQSIDg1XEdru3aNt0fPOjberh5dG2/P/f414+2ez/5uvN3+n2XOB+HU5AUJEAQSIAgkQBBIgCCQAEEgAYJAAgSBBAgCCRAEEiA4NVxHxw8fHm/3XjXfrpfjG/0F4CTjBQkQBBIgCCRAEEiAIJAAQSABgkACBIEECAIJEAQSIAgkQBBIgCCQAEEgAYJAAgSBBAgCCRAEEiAIJEAQSIAgkABBIAGCQAIEgQQIAgkQBBIgCCRAEEiAIJAAQSABgkACBIEECAIJEAQSIAgkQBBIgCCQAEEgAYJAAgSBBAgCCRAEEiAIJEAQSIAgkABBIAGCQAIEgQQIAgkQBBIgCCRAEEiAIJAAQSABgkACBIEECAIJEAQSIAgkQBBIgCCQAEEgAYJAAgSBBAgCCRBWF4vFYqO/BMDJyAsSIAgkQBBIgCCQAEEgAYJAAgSBBAgCCRAEEiD8D7AYv4rYgrpbAAAAAElFTkSuQmCC\n"
          },
          "metadata": {}
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "<Figure size 400x400 with 16 Axes>"
            ],
            "image/png": "iVBORw0KGgoAAAANSUhEUgAAAUkAAAFICAYAAADd1gwNAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjcuMSwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy/bCgiHAAAACXBIWXMAAA9hAAAPYQGoP6dpAAAIIUlEQVR4nO3dX6jfdR3H8e/5k4sNLdcxobB2zJ1GeCFpxawIZstRF14tsrKMRcKCRAtLIshEKQJbUZFREpWwWDcZ2cVxzILmqdysJUwlZ5mC9seicqLn7Hy7HMzx+rz7nfPb+bPH4/r9e58v33POk8/F98tvpO/7vgPgpEaX+gIAljORBAhEEiAQSYBAJAECkQQIRBIgEEmAQCQBgvHq4NbR7cO8jhVlen7PQJ9zD49zDxfOPVy4yj10kgQIRBIgEEmAQCQBApEECEQSIBBJgEAkAQKRBAhEEiAQSYBAJAECkQQIRBIgEEmAQCQBApEECEQSIBBJgEAkAQKRBAhEEiAQSYBAJAECkQQIRBIgEEmAQCQBApEECEQSIBBJgEAkAQKRBAhEEiAQSYBAJAECkQQIRBIgEEmAQCQBApEECEQSIBBJgEAkAQKRBAhEEiAQSYBAJAECkQQIRBIgEEmAQCQBApEECEQSIBBJgEAkAQKRBAhEEiAQSYBAJAECkQQIRBIgEEmAQCQBApEECEQSIBjp+75f6osAWK6cJAECkQQIRBIgEEmAQCQBApEECEQSIBBJgEAkAQKRBAhEEiAQSYBgvDq4dXT7MK9jRZme3zPQ59zD49zDhXMPF65yD50kAQKRBAhEEiAQSYBAJAECkQQIRBIgEEmAQCQBgvIbN8DpZWziFYu265nLNzZnnn77fG3ZmmPNkamPHKjtKnCSBAhEEiAQSYBAJAECkQQIRBIgEEmAQCQBAg+Ts6LMbbm4NPeXj802Z177jbHSrmdftaY0V/HMG9rnkufPnVu0n3cyZ/9qfWlu9+TeRfuZu/75aHPma3svL+162aHF+31UOEkCBCIJEIgkQCCSAIFIAgQiCRCIJEAgkgCBSAIE3rhhUT117aVD3b/3h98tze14/G3Nmau+t7+0a/3Y0ebMXf++qLSr4o7fDfcePnDv60tzb93V/sqFMx97trSrv//B5szG7telXaeakyRAIJIAgUgCBCIJEIgkQCCSAIFIAgQiCRB4mJyyJ25sP+T8i51fLm67bqBrOP/H15TmNn39782ZW175odKu8YOPNGfmj7YfOK/a2B2sDX5wsP0bPnvfYB88iX7RNi1fTpIAgUgCBCIJEIgkQCCSAIFIAgQiCRCIJEAgkgCBN25WufENr2nOHP7CRGnXoS23NWcu2v3J0q4j15fGXmTq+gOluWNzc82Z0faLNF3Xdd18bYxVykkSIBBJgEAkAQKRBAhEEiAQSYBAJAECkQQIPEy+Qj13xZtLc++6+ZfNmZ9N3FXa9cb7P9yc6cdKqwb2nt//rTS3fvy/zZnPzVxR2jX10QebM/3sC6VdrDxOkgCBSAIEIgkQiCRAIJIAgUgCBCIJEIgkQCCSAMFI3/f9Ul8EwHLlJAkQiCRAIJIAgUgCBCIJEIgkQCCSAIFIAgQiCRCIJEAgkgCBSAIEIgkQlL93e+vo9mFex4oyPb9noM9tO3dnae7wrRuaMwe3fbW066zRlzZn/vDCbGnXrU++uznz28Pnl3b9eccNpbkTXfCl20pzL3+oPfPtz+8q7Xr/d65rzpx3y/7SrsU06N+h/+XjKvfQSRIgEEmAQCQBApEECEQSIBBJgEAkAQKRBAhEEiAov3HDwr3lnidLc3efM92cOfB87Vd35Y8+3pyZ/Mx9pV1d94/mxFRhpuu6rttR/JEnmLyxdq2ja9e2Z26qfeX87Fm+mv505iQJEIgkQCCSAIFIAgQiCRCIJEAgkgCBSAIEHiY/hfZfc0lp7rJ1m5szZ/zmkdKuyf9UHxRfGUbXrSvNPfzFC5szs/29pV2TPzlammN1cpIECEQSIBBJgEAkAQKRBAhEEiAQSYBAJAECkQQIvHFzKs0cKo1VfinzC7uSoRlZs2ao+x///obS3JHNtzdnLr7p2tKuif2r660l/j9OkgCBSAIEIgkQiCRAIJIAgUgCBCIJEIgkQOBhcspGLml/JcKm2x8a6jXc86b2Q+Jd13Xv/MAnmjMT+zwkTpuTJEAgkgCBSAIEIgkQiCRAIJIAgUgCBCIJEIgkQOCNm1Vu5CVnNGeO/GBTadfPN3+zOfOVv15W2jWoq9+7szQ3NnNwqNfB6cNJEiAQSYBAJAECkQQIRBIgEEmAQCQBApEECDxMvgz96ebNzZljr3uutOvhd9zRnHl0bl9p11Wf/lRz5szdM6Vd3Xxt7EVmDg34QRiMkyRAIJIAgUgCBCIJEIgkQCCSAIFIAgQiCRCIJEAw0vd9v9QXAbBcOUkCBCIJEIgkQCCSAIFIAgQiCRCIJEAgkgCBSAIEIgkQiCRAIJIAQfkrZbeObh/mdawo0/N7BvrctqkbSnN37ruzOXP22NrSrm/969XNmZ9uubC0a+6pp0tzFYPeQ3+Hx7mHC1e5h06SAIFIAgQiCRCIJEAgkgCBSAIEIgkQiCRAIJIAQfmNGxbu2B8fK82977xLh3wlJ1q8N2lgtXGSBAhEEiAQSYBAJAECkQQIRBIgEEmAQCQBApEECEQSIBBJgEAkAQKRBAhEEiAQSYBAJAECkQQIRBIgEEmAQCQBApEECEQSIBBJgEAkAQKRBAhEEiAQSYBAJAECkQQIRBIgEEmAQCQBApEECEQSIBBJgEAkAQKRBAhEEiAQSYBAJAECkQQIRBIgEEmAQCQBApEECEQSIBBJgEAkAQKRBAhEEiAQSYBAJAECkQQIRBIgEEmAQCQBgpG+7/ulvgiA5cpJEiAQSYBAJAECkQQIRBIgEEmAQCQBApEECEQSIPgfGaf2A6XA0XAAAAAASUVORK5CYII=\n"
          },
          "metadata": {}
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "Running the experiment"
      ],
      "metadata": {
        "id": "87RZchNxkujB"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "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=keras.losses.SparseCategoricalCrossentropy(from_logits=True),\n",
        "        metrics=[\n",
        "            keras.metrics.SparseCategoricalAccuracy(name=\"accuracy\"),\n",
        "            keras.metrics.SparseTopKCategoricalAccuracy(5, name=\"top-5-accuracy\"),\n",
        "        ],\n",
        "    )\n",
        "\n",
        "    checkpoint_filepath = \"/tmp/checkpoint\"\n",
        "    checkpoint_callback = 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",
        "        x=x_train,\n",
        "        y=y_train,\n",
        "        batch_size=batch_size,\n",
        "        epochs=num_epochs,\n",
        "        validation_split=0.1,\n",
        "        callbacks=[checkpoint_callback],\n",
        "    )\n",
        "\n",
        "    model.load_weights(checkpoint_filepath)\n",
        "    _, accuracy, top_5_accuracy = model.evaluate(x_test, y_test)\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",
        "\n",
        "vit_classifier = create_vit_classifier()\n",
        "history = run_experiment(vit_classifier)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "A5bP6SxHktDg",
        "outputId": "c0c6737d-3651-41d2-ae1a-d4a5362a7073"
      },
      "execution_count": 57,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Epoch 1/5\n",
            "27/27 [==============================] - 846s 31s/step - loss: 1.4018 - accuracy: 0.6078 - top-5-accuracy: 0.9006 - val_loss: 0.2222 - val_accuracy: 0.9280 - val_top-5-accuracy: 0.9960\n",
            "Epoch 2/5\n",
            "27/27 [==============================] - 814s 30s/step - loss: 0.3816 - accuracy: 0.8777 - top-5-accuracy: 0.9938 - val_loss: 0.1178 - val_accuracy: 0.9652 - val_top-5-accuracy: 0.9975\n",
            "Epoch 3/5\n",
            "27/27 [==============================] - 804s 30s/step - loss: 0.2583 - accuracy: 0.9200 - top-5-accuracy: 0.9966 - val_loss: 0.0879 - val_accuracy: 0.9748 - val_top-5-accuracy: 0.9982\n",
            "Epoch 4/5\n",
            "27/27 [==============================] - 846s 31s/step - loss: 0.1912 - accuracy: 0.9407 - top-5-accuracy: 0.9981 - val_loss: 0.0724 - val_accuracy: 0.9765 - val_top-5-accuracy: 0.9985\n",
            "Epoch 5/5\n",
            "27/27 [==============================] - 815s 30s/step - loss: 0.1605 - accuracy: 0.9501 - top-5-accuracy: 0.9985 - val_loss: 0.0627 - val_accuracy: 0.9825 - val_top-5-accuracy: 0.9985\n",
            "313/313 [==============================] - 46s 140ms/step - loss: 0.0689 - accuracy: 0.9767 - top-5-accuracy: 0.9996\n",
            "Test accuracy: 97.67%\n",
            "Test top 5 accuracy: 99.96%\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "vit_classifier.evaluate(x_test, y_test)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "TWOiWL8Ly01p",
        "outputId": "3759b227-345a-48af-8285-acd5cb049836"
      },
      "execution_count": 64,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "313/313 [==============================] - 54s 173ms/step - loss: 0.0689 - accuracy: 0.9767 - top-5-accuracy: 0.9996\n"
          ]
        },
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "[0.06887061893939972, 0.9767000079154968, 0.9995999932289124]"
            ]
          },
          "metadata": {},
          "execution_count": 64
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "prediction = vit_classifier.predict(x_test[:,:,:,0])"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "CGZGH2-r6S4F",
        "outputId": "e2ebc595-9cfc-4966-9235-068f0d364ac8"
      },
      "execution_count": 75,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "313/313 [==============================] - 54s 174ms/step\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "pred_img_num = 199\n",
        "image = np.squeeze(x_test[pred_img_num,:,:,0])\n",
        "plt.imshow(image)\n",
        "plt.show()\n",
        "\n",
        "print(\"Predicted:\", np.argmax(prediction[pred_img_num,:]))"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 447
        },
        "id": "t17jNGLV73xh",
        "outputId": "197a7061-23ba-4323-c243-abdd81409998"
      },
      "execution_count": 90,
      "outputs": [
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "<Figure size 640x480 with 1 Axes>"
            ],
            "image/png": "iVBORw0KGgoAAAANSUhEUgAAAaAAAAGdCAYAAABU0qcqAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjcuMSwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy/bCgiHAAAACXBIWXMAAA9hAAAPYQGoP6dpAAAbn0lEQVR4nO3df3DUdZ7n8VcHSAOadCbGpBMJTACVUSDeIGSyKoNDihB3WX7VFP6YLbBcLJngChlHK3MqOmNdZuBWPd0ot3szMEyJqFcCJeUwhcGEc0xwQViWUXOEi0MYSFBu6Q5BQiCf+4OztSWA36Y773R4Pqq+VaT7+8n3zdcun3zp5hufc84JAIBelmI9AADg8kSAAAAmCBAAwAQBAgCYIEAAABMECABgggABAEwQIACAiYHWA3xdd3e3Dh06pLS0NPl8PutxAAAeOefU3t6uvLw8paSc/zqnzwXo0KFDys/Ptx4DAHCJWlpaNGzYsPM+3+cClJaWJkm6VXdooAYZTwMA8Oq0uvSu3or8//x8Ehag6upqrVixQq2trSosLNQLL7ygSZMmXXTdF3/tNlCDNNBHgAAg6fz/O4xe7G2UhHwI4dVXX1VFRYWWLVumDz74QIWFhSotLdWRI0cScTgAQBJKSICeeeYZLVy4UPfee69uuOEGrVy5UkOHDtVvfvObRBwOAJCE4h6gU6dOaefOnSopKfnyICkpKikpUX19/Tn7d3Z2KhwOR20AgP4v7gH67LPPdObMGeXk5EQ9npOTo9bW1nP2r6qqUiAQiGx8Ag4ALg/m/xC1srJSoVAosrW0tFiPBADoBXH/FFxWVpYGDBigtra2qMfb2toUDAbP2d/v98vv98d7DABAHxf3K6DU1FRNmDBBNTU1kce6u7tVU1Oj4uLieB8OAJCkEvLvgCoqKjR//nzdfPPNmjRpkp577jl1dHTo3nvvTcThAABJKCEBmjdvnj799FM98cQTam1t1U033aTNmzef88EEAMDly+ecc9ZDfFU4HFYgENAUzeROCACQhE67LtVqo0KhkNLT08+7n/mn4AAAlycCBAAwQYAAACYIEADABAECAJggQAAAEwQIAGCCAAEATBAgAIAJAgQAMEGAAAAmCBAAwAQBAgCYIEAAABMECABgggABAEwQIACACQIEADBBgAAAJggQAMAEAQIAmCBAAAATBAgAYIIAAQBMECAAgAkCBAAwQYAAACYIEADABAECAJggQAAAEwQIAGCCAAEATBAgAIAJAgQAMEGAAAAmCBAAwAQBAgCYIEAAABMECABgYqD1AEBf0nnHRM9rDpR5/3Pc/5n73z2v6XJnPK/506nTntdI0pxN/+B5zZBDAzyvyV/xvuc17nRsvyf0PVwBAQBMECAAgAkCBAAwQYAAACYIEADABAECAJggQAAAEwQIAGCCAAEATBAgAIAJAgQAMEGAAAAmuBkp+ryBw67xvOY/bs2P6VijH/rQ85o/DK/xvKbLef+zX7e6Pa/5Tmpsf8b8aM4/xbTOqxuyH/S8ZvTShgRMAgtcAQEATBAgAICJuAfoySeflM/ni9rGjBkT78MAAJJcQt4DuvHGG/X2229/eZCBvNUEAIiWkDIMHDhQwWAwEd8aANBPJOQ9oH379ikvL08jR47UPffcowMHDpx3387OToXD4agNAND/xT1ARUVFWr16tTZv3qyXXnpJzc3Nuu2229Te3t7j/lVVVQoEApEtPz+2j88CAJJL3ANUVlamH/7whxo/frxKS0v11ltv6dixY3rttdd63L+yslKhUCiytbS0xHskAEAflPBPB2RkZOi6665TU1NTj8/7/X75/f5EjwEA6GMS/u+Ajh8/rv379ys3NzfRhwIAJJG4B+jhhx9WXV2dPvnkE7333nuaPXu2BgwYoLvuuivehwIAJLG4/xXcwYMHddddd+no0aO6+uqrdeutt6qhoUFXX311vA8FAEhicQ/QunXr4v0t0Y9033qT5zXLfvcvntcUpnpeAgPP/s0az2teeGOe5zUp/2uX5zVIPO4FBwAwQYAAACYIEADABAECAJggQAAAEwQIAGCCAAEATBAgAIAJAgQAMEGAAAAmCBAAwAQBAgCYSPgPpAO+6vjwwZ7XcGPR/qt0aMjzmt88fcjzms6yoZ7XSFL3iRMxrcM3wxUQAMAEAQIAmCBAAAATBAgAYIIAAQBMECAAgAkCBAAwQYAAACYIEADABAECAJggQAAAEwQIAGCCAAEATHA3bOASLT44xfOavc+Ni/8gPejIi+3PmP9a8d/iPEn8vDp6k+c1s66bH9vBdn8Y2zp8I1wBAQBMECAAgAkCBAAwQYAAACYIEADABAECAJggQAAAEwQIAGCCAAEATBAgAIAJAgQAMEGAAAAmuBkpelXm1mbPa25+7iHPa64p+7PnNZK079/yPa+59qc7PK9JP93geU0sAv/pxtgWVsR3DqAnXAEBAEwQIACACQIEADBBgAAAJggQAMAEAQIAmCBAAAATBAgAYIIAAQBMECAAgAkCBAAwQYAAACa4GSl61enWNs9r8lZ4X+NWeF4iSRqtv3g/VmyHAi57XAEBAEwQIACACc8B2rZtm2bMmKG8vDz5fD5t2LAh6nnnnJ544gnl5uZqyJAhKikp0b59++I1LwCgn/AcoI6ODhUWFqq6urrH55cvX67nn39eK1eu1Pbt23XFFVeotLRUJ0+evORhAQD9h+cPIZSVlamsrKzH55xzeu655/TYY49p5syZkqQ1a9YoJydHGzZs0J133nlp0wIA+o24vgfU3Nys1tZWlZSURB4LBAIqKipSfX19j2s6OzsVDoejNgBA/xfXALW2tkqScnJyoh7PycmJPPd1VVVVCgQCkS0/Pz+eIwEA+ijzT8FVVlYqFApFtpaWFuuRAAC9IK4BCgaDkqS2tuh/ONjW1hZ57uv8fr/S09OjNgBA/xfXABUUFCgYDKqmpibyWDgc1vbt21VcXBzPQwEAkpznT8EdP35cTU1Nka+bm5u1e/duZWZmavjw4VqyZImefvppXXvttSooKNDjjz+uvLw8zZo1K55zAwCSnOcA7dixQ7fffnvk64qKCknS/PnztXr1aj3yyCPq6OjQ/fffr2PHjunWW2/V5s2bNXjw4PhNDQBIej7nXJ+6l2I4HFYgENAUzdRA3yDrcYCk9r//eWJM6z7+6xfjPEn8/O3Hsz2v8d3xaUzHcp2dMa273J12XarVRoVCoQu+r2/+KTgAwOWJAAEATBAgAIAJAgQAMEGAAAAmCBAAwAQBAgCYIEAAABMECABgggABAEwQIACACQIEADBBgAAAJjz/OAYANvY9X+R5Td30f4zxaP4Y13mz/ni25zW+RwKe17jOg57XIPG4AgIAmCBAAAATBAgAYIIAAQBMECAAgAkCBAAwQYAAACYIEADABAECAJggQAAAEwQIAGCCAAEATHAzUsBA+7zveV7z/mzvNxZNS+mdm4rGavVf/srzGrfzTwmYBBa4AgIAmCBAAAATBAgAYIIAAQBMECAAgAkCBAAwQYAAACYIEADABAECAJggQAAAEwQIAGCCAAEATHAzUuASxXJj0feeXel5TZcb7HlNb1oVzve8JmXREM9rznhegb6KKyAAgAkCBAAwQYAAACYIEADABAECAJggQAAAEwQIAGCCAAEATBAgAIAJAgQAMEGAAAAmCBAAwAQ3IwW+IvQj7zcW3VT1j57XxHJj0W51e14Tq/buU57X/O4/z/C8Zmjjds9r0H9wBQQAMEGAAAAmPAdo27ZtmjFjhvLy8uTz+bRhw4ao5xcsWCCfzxe1TZ8+PV7zAgD6Cc8B6ujoUGFhoaqrq8+7z/Tp03X48OHI9sorr1zSkACA/sfzhxDKyspUVlZ2wX38fr+CwWDMQwEA+r+EvAdUW1ur7OxsXX/99Vq0aJGOHj163n07OzsVDoejNgBA/xf3AE2fPl1r1qxRTU2NfvWrX6murk5lZWU6c6bnn+ReVVWlQCAQ2fLzvf9ceQBA8on7vwO68847I78eN26cxo8fr1GjRqm2tlZTp049Z//KykpVVFREvg6Hw0QIAC4DCf8Y9siRI5WVlaWmpqYen/f7/UpPT4/aAAD9X8IDdPDgQR09elS5ubmJPhQAIIl4/iu448ePR13NNDc3a/fu3crMzFRmZqaeeuopzZ07V8FgUPv379cjjzyi0aNHq7S0NK6DAwCSm+cA7dixQ7fffnvk6y/ev5k/f75eeukl7dmzR7/97W917Ngx5eXladq0afrFL34hv98fv6kBAEnPc4CmTJki59x5n//DH/5wSQMB8fAf84tjWvc/f7HC85q0lP73h6s1oXGe16Tt/czzmp4/G4vLBfeCAwCYIEAAABMECABgggABAEwQIACACQIEADBBgAAAJggQAMAEAQIAmCBAAAATBAgAYIIAAQBMECAAgIm4/0huIN7+773e72y97knvd7WWpJwB/e/O1rEo/1aj5zVFm3v+qccX8ne/X+R5zZjKjzyvORMOe16DxOMKCABgggABAEwQIACACQIEADBBgAAAJggQAMAEAQIAmCBAAAATBAgAYIIAAQBMECAAgAkCBAAwwc1IETP3V4We17QVXeF5zbqH/qvnNcMGclPR3naz/4znNR/N+ifPayZv/wfPazLW1Hteg8TjCggAYIIAAQBMECAAgAkCBAAwQYAAACYIEADABAECAJggQAAAEwQIAGCCAAEATBAgAIAJAgQAMMHNSPuZARkBz2v2vVgQ07HWFf+z5zVjU30xHGlQDGuQDLZ3ev9v6w95v+kp+iaugAAAJggQAMAEAQIAmCBAAAATBAgAYIIAAQBMECAAgAkCBAAwQYAAACYIEADABAECAJggQAAAE9yMtJ/pLhjmec3e7/9LjEeL5caiwJfmb1noec11G99PwCSwwBUQAMAEAQIAmPAUoKqqKk2cOFFpaWnKzs7WrFmz1NjYGLXPyZMnVV5erquuukpXXnml5s6dq7a2trgODQBIfp4CVFdXp/LycjU0NGjLli3q6urStGnT1NHREdln6dKlevPNN/X666+rrq5Ohw4d0pw5c+I+OAAguXn6EMLmzZujvl69erWys7O1c+dOTZ48WaFQSL/+9a+1du1a/eAHP5AkrVq1St/5znfU0NCg733ve/GbHACQ1C7pPaBQKCRJyszMlCTt3LlTXV1dKikpiewzZswYDR8+XPX19T1+j87OToXD4agNAND/xRyg7u5uLVmyRLfccovGjh0rSWptbVVqaqoyMjKi9s3JyVFra2uP36eqqkqBQCCy5efnxzoSACCJxByg8vJy7d27V+vWrbukASorKxUKhSJbS0vLJX0/AEByiOkfoi5evFibNm3Stm3bNGzYl//wMRgM6tSpUzp27FjUVVBbW5uCwWCP38vv98vv98cyBgAgiXm6AnLOafHixVq/fr22bt2qgoKCqOcnTJigQYMGqaamJvJYY2OjDhw4oOLi4vhMDADoFzxdAZWXl2vt2rXauHGj0tLSIu/rBAIBDRkyRIFAQPfdd58qKiqUmZmp9PR0PfjggyouLuYTcACAKJ4C9NJLL0mSpkyZEvX4qlWrtGDBAknSs88+q5SUFM2dO1ednZ0qLS3Viy++GJdhAQD9h6cAOecuus/gwYNVXV2t6urqmIdC7BoXDbUeAQnSdqbT85qS936cgEl6duaQ99femGV/8rym2/MK9FXcCw4AYIIAAQBMECAAgAkCBAAwQYAAACYIEADABAECAJggQAAAEwQIAGCCAAEATBAgAIAJAgQAMEGAAAAmYvqJqOi7hn4yyHqEpPZfPrvJ85o1O7z/sMUbnv7U8xqdPuN5SUHLHu/H6UXc2fryxhUQAMAEAQIAmCBAAAATBAgAYIIAAQBMECAAgAkCBAAwQYAAACYIEADABAECAJggQAAAEwQIAGCCm5H2M9/+7See14wb9/cxHevfJ/+PmNZ5NW6b9/nSaofGdKzs3/2b5zXXndjhec1pzyuA/ocrIACACQIEADBBgAAAJggQAMAEAQIAmCBAAAATBAgAYIIAAQBMECAAgAkCBAAwQYAAACYIEADABDcj7WdO/+WQ5zUFd3lfI0l/q4kxrfOqQN5vEBqr7l47EgCugAAAJggQAMAEAQIAmCBAAAATBAgAYIIAAQBMECAAgAkCBAAwQYAAACYIEADABAECAJggQAAAEwQIAGCCAAEATBAgAIAJTwGqqqrSxIkTlZaWpuzsbM2aNUuNjY1R+0yZMkU+ny9qe+CBB+I6NAAg+XkKUF1dncrLy9XQ0KAtW7aoq6tL06ZNU0dHR9R+Cxcu1OHDhyPb8uXL4zo0ACD5efqJqJs3b476evXq1crOztbOnTs1efLkyONDhw5VMBiMz4QAgH7pkt4DCoVCkqTMzMyox19++WVlZWVp7Nixqqys1IkTJ877PTo7OxUOh6M2AED/5+kK6Ku6u7u1ZMkS3XLLLRo7dmzk8bvvvlsjRoxQXl6e9uzZo0cffVSNjY164403evw+VVVVeuqpp2IdAwCQpHzOORfLwkWLFun3v/+93n33XQ0bNuy8+23dulVTp05VU1OTRo0adc7znZ2d6uzsjHwdDoeVn5+vKZqpgb5BsYwGADB02nWpVhsVCoWUnp5+3v1iugJavHixNm3apG3btl0wPpJUVFQkSecNkN/vl9/vj2UMAEAS8xQg55wefPBBrV+/XrW1tSooKLjomt27d0uScnNzYxoQANA/eQpQeXm51q5dq40bNyotLU2tra2SpEAgoCFDhmj//v1au3at7rjjDl111VXas2ePli5dqsmTJ2v8+PEJ+Q0AAJKTp/eAfD5fj4+vWrVKCxYsUEtLi370ox9p79696ujoUH5+vmbPnq3HHnvsgn8P+FXhcFiBQID3gAAgSSXkPaCLtSo/P191dXVeviUA4DLFveAAACYIEADABAECAJggQAAAEwQIAGCCAAEATBAgAIAJAgQAMEGAAAAmCBAAwAQBAgCYIEAAABMECABgggABAEwQIACACQIEADBBgAAAJggQAMAEAQIAmCBAAAATBAgAYIIAAQBMECAAgAkCBAAwQYAAACYGWg/wdc45SdJpdUnOeBgAgGen1SXpy/+fn0+fC1B7e7sk6V29ZTwJAOBStLe3KxAInPd5n7tYonpZd3e3Dh06pLS0NPl8vqjnwuGw8vPz1dLSovT0dKMJ7XEezuI8nMV5OIvzcFZfOA/OObW3tysvL08pKed/p6fPXQGlpKRo2LBhF9wnPT39sn6BfYHzcBbn4SzOw1mch7Osz8OFrny+wIcQAAAmCBAAwERSBcjv92vZsmXy+/3Wo5jiPJzFeTiL83AW5+GsZDoPfe5DCACAy0NSXQEBAPoPAgQAMEGAAAAmCBAAwETSBKi6ulrf/va3NXjwYBUVFen999+3HqnXPfnkk/L5fFHbmDFjrMdKuG3btmnGjBnKy8uTz+fThg0bop53zumJJ55Qbm6uhgwZopKSEu3bt89m2AS62HlYsGDBOa+P6dOn2wybIFVVVZo4caLS0tKUnZ2tWbNmqbGxMWqfkydPqry8XFdddZWuvPJKzZ07V21tbUYTJ8Y3OQ9Tpkw55/XwwAMPGE3cs6QI0KuvvqqKigotW7ZMH3zwgQoLC1VaWqojR45Yj9brbrzxRh0+fDiyvfvuu9YjJVxHR4cKCwtVXV3d4/PLly/X888/r5UrV2r79u264oorVFpaqpMnT/bypIl1sfMgSdOnT496fbzyyiu9OGHi1dXVqby8XA0NDdqyZYu6uro0bdo0dXR0RPZZunSp3nzzTb3++uuqq6vToUOHNGfOHMOp4++bnAdJWrhwYdTrYfny5UYTn4dLApMmTXLl5eWRr8+cOePy8vJcVVWV4VS9b9myZa6wsNB6DFOS3Pr16yNfd3d3u2Aw6FasWBF57NixY87v97tXXnnFYMLe8fXz4Jxz8+fPdzNnzjSZx8qRI0ecJFdXV+ecO/vfftCgQe7111+P7PPRRx85Sa6+vt5qzIT7+nlwzrnvf//77qGHHrIb6hvo81dAp06d0s6dO1VSUhJ5LCUlRSUlJaqvrzeczMa+ffuUl5enkSNH6p577tGBAwesRzLV3Nys1tbWqNdHIBBQUVHRZfn6qK2tVXZ2tq6//notWrRIR48etR4poUKhkCQpMzNTkrRz5051dXVFvR7GjBmj4cOH9+vXw9fPwxdefvllZWVlaezYsaqsrNSJEycsxjuvPncz0q/77LPPdObMGeXk5EQ9npOTo48//thoKhtFRUVavXq1rr/+eh0+fFhPPfWUbrvtNu3du1dpaWnW45lobW2VpB5fH188d7mYPn265syZo4KCAu3fv18/+9nPVFZWpvr6eg0YMMB6vLjr7u7WkiVLdMstt2js2LGSzr4eUlNTlZGREbVvf3499HQeJOnuu+/WiBEjlJeXpz179ujRRx9VY2Oj3njjDcNpo/X5AOFLZWVlkV+PHz9eRUVFGjFihF577TXdd999hpOhL7jzzjsjvx43bpzGjx+vUaNGqba2VlOnTjWcLDHKy8u1d+/ey+J90As533m4//77I78eN26ccnNzNXXqVO3fv1+jRo3q7TF71Of/Ci4rK0sDBgw451MsbW1tCgaDRlP1DRkZGbruuuvU1NRkPYqZL14DvD7ONXLkSGVlZfXL18fixYu1adMmvfPOO1E/viUYDOrUqVM6duxY1P799fVwvvPQk6KiIknqU6+HPh+g1NRUTZgwQTU1NZHHuru7VVNTo+LiYsPJ7B0/flz79+9Xbm6u9ShmCgoKFAwGo14f4XBY27dvv+xfHwcPHtTRo0f71evDOafFixdr/fr12rp1qwoKCqKenzBhggYNGhT1emhsbNSBAwf61evhYuehJ7t375akvvV6sP4UxDexbt065/f73erVq92HH37o7r//fpeRkeFaW1utR+tVP/nJT1xtba1rbm52f/zjH11JSYnLyspyR44csR4todrb292uXbvcrl27nCT3zDPPuF27drk///nPzjnnfvnLX7qMjAy3ceNGt2fPHjdz5kxXUFDgPv/8c+PJ4+tC56G9vd09/PDDrr6+3jU3N7u3337bffe733XXXnutO3nypPXocbNo0SIXCARcbW2tO3z4cGQ7ceJEZJ8HHnjADR8+3G3dutXt2LHDFRcXu+LiYsOp4+9i56Gpqcn9/Oc/dzt27HDNzc1u48aNbuTIkW7y5MnGk0dLigA559wLL7zghg8f7lJTU92kSZNcQ0OD9Ui9bt68eS43N9elpqa6a665xs2bN881NTVZj5Vw77zzjpN0zjZ//nzn3NmPYj/++OMuJyfH+f1+N3XqVNfY2Gg7dAJc6DycOHHCTZs2zV199dVu0KBBbsSIEW7hwoX97g9pPf3+JblVq1ZF9vn888/dj3/8Y/etb33LDR061M2ePdsdPnzYbugEuNh5OHDggJs8ebLLzMx0fr/fjR492v30pz91oVDIdvCv4ccxAABM9Pn3gAAA/RMBAgCYIEAAABMECABgggABAEwQIACACQIEADBBgAAAJggQAMAEAQIAmCBAAAATBAgAYOL/AUznpgdE8ndEAAAAAElFTkSuQmCC\n"
          },
          "metadata": {}
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Predicted: 2\n"
          ]
        }
      ]
    }
  ]
}