{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"pygments_lexer":"ipython3","nbconvert_exporter":"python","version":"3.6.4","file_extension":".py","codemirror_mode":{"name":"ipython","version":3},"name":"python","mimetype":"text/x-python"}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"![image](https://user-images.githubusercontent.com/17668390/149625554-b9c7074a-2137-49d5-8726-a3fbfa3f9a4c.gif)\n\n<div align=\"center\">\n    Figure: <strong>Grad-CAM</strong> of Hybrid-EfficientNet-Swin Transformer. Left (<i>Input</i>), Middle (<i>EfficientNet</i>), Right (<i>Swin Transformer</i>).\n</div>","metadata":{"id":"m3nFPpkKOJMR"}},{"cell_type":"markdown","source":"<table class=\"tfo-notebook-buttons\" align=\"center\">\n    \n  <td>\n    <a target=\"_blank\" href=\"https://colab.research.google.com/drive/1usxq9yhZthyapAnzFfFObQ7RjXPicAop?usp=sharing\"><img src=\"https://www.tensorflow.org/images/colab_logo_32px.png\" />Run in Google Colab</a>\n  </td>\n    \n  <td>\n    <a target=\"_blank\" href=\"https://github.com/innat/HybridModel-GradCAM\"><img src=\"https://www.tensorflow.org/images/GitHub-Mark-32px.png\" />View source on GitHub</a>\n  </td>\n    \n   <td>\n    <a target=\"_blank\" href=\"https://deepnote.com/workspace/mohammed-innat-36e929bc-ce23-4d95-9ddc-a9c6662eb7d6/project/Notebooks-32e94ef4-8ce0-4cc4-8042-49862519f3f2/%2F%5BDeepnote%5D_HENetSwinT.ipynb\"><img src=\"https://user-images.githubusercontent.com/17668390/176064308-845cc64d-cd84-44fa-a491-4a53759b19d4.png\" />Run in Deepnote</a>\n  </td>\n    \n   <td>\n    <a target=\"_blank\" href=\"https://www.kaggle.com/code/ipythonx/tf-hybrid-efficientnet-swin-transformer-gradcam\"><img src=\"https://user-images.githubusercontent.com/17668390/176064379-9bcc7836-dcff-42b2-bcca-feff28f22c70.png\" />Run in Kaggle</a>\n   </td>\n    \n   <td>\n    <a target=\"_blank\" href=\"https://huggingface.co/spaces/innat/HybridModel-GradCAM\"><img src=\"https://user-images.githubusercontent.com/17668390/176064420-46cbf547-0d17-4438-a791-d23e17eff5a9.png\" />Try on Gradio</a>\n   </td>\n    \n</table>\n\n# Introduction\n\nConvolutional Neural Networks (CNNs) have been the de facto model for visual data up to this point. But the Vision transformers have proven as an efficient alternative for CNNs. Current research has shown that transformer models can perform similarly, if not better, on vison tasks as well. In this code example, we have tried to inspect the visual interpretations of a **CNN** and **Transformer** blocks of a hybrid model (EfficientNet + Swin Transformer) with the **GradCAM** technique. In the result, it appears that the transformer blocks are capable of globally refining feature activation across the relevant object, as opposed to the CNN, which is more focused on operating locally (shown in above figure). However, the approach that will be shown here, is highly experimental. The workflow probably can generate a more meaningful modeling approach.","metadata":{"id":"d8NkI9z0OJMa"}},{"cell_type":"markdown","source":"**Data Pipelines**. To keep things simple, we've used [tf_flowers](https://www.tensorflow.org/datasets/catalog/tf_flowers) dataset, a multi-class classification problem. To apply image augmentaiton, we have used the **vectorized** implementaiton of **CutMix** and **MixUp**, derived from the [KerasCV](https://keras.io/keras_cv/). Additionally, we have also used [Jax](https://jax.readthedocs.io/en/latest/) library to write image augmentaiton layer and used it in the `tf.data` API. FYI, yes we can use `Jax` code to build `keras` layer. Also, we have replaced all possible `numpy` code to `jax.numpy` wherever possible.\n\n\n- [**Code Style: Keras.io.**](https://bit.ly/3Oe9zHY) ✔️\n- **JIT Compilation** ✔️\n- **Mixed Precision** ✔️\n- **Gradient Accumulation** ✔️\n- **Label Smoothing** ✔️\n- **TensorFlow Lite Conversion** ✔️\n\n\n**Note**: If the accelerator is set to **GPU**, training will be the first to begin, followed by inference. However, if the accelerator is set to **CPU**, only inference will be performed, and a trained weight file will be used. Also, this notebook is tested on `Jax 0.3.13` and `TensorFlow 2.6.4` on Kaggle, `TF 2.8` on Colab, and `TF 2.9` on Deepnote.","metadata":{"id":"bQ-gOUg4OJMc"}},{"cell_type":"code","source":"!pip install tensorflow==2.9.2","metadata":{"execution":{"iopub.status.busy":"2023-11-08T04:11:20.375647Z","iopub.execute_input":"2023-11-08T04:11:20.375983Z","iopub.status.idle":"2023-11-08T04:12:40.755957Z","shell.execute_reply.started":"2023-11-08T04:11:20.375956Z","shell.execute_reply":"2023-11-08T04:12:40.753289Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"!pip install gdown -q\n!pip install jax==0.3.13 jaxlib==0.3.10 -q\n# !pip install tensorflow==2.9.2","metadata":{"id":"T-u93ThSOJMd","_kg_hide-output":true,"execution":{"iopub.status.busy":"2023-11-08T04:12:40.760561Z","iopub.execute_input":"2023-11-08T04:12:40.761166Z","iopub.status.idle":"2023-11-08T04:12:52.952637Z","shell.execute_reply.started":"2023-11-08T04:12:40.761126Z","shell.execute_reply":"2023-11-08T04:12:52.951869Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import os\nimport random\nimport warnings\n\nimport cv2\nimport gdown\nfrom functools import partial\n\nwarnings.simplefilter(action=\"ignore\")\nos.environ[\"TF_CPP_MIN_LOG_LEVEL\"] = \"3\"\nos.environ[\"XLA_PYTHON_CLIENT_PREALLOCATE\"] = \"false\"\n\nimport numpy as np\nimport pandas as pd\nfrom matplotlib import cm\nfrom numpy.random import rand\nimport matplotlib.pyplot as plt\n\nimport tensorflow as tf\nfrom tensorflow import keras\nfrom tensorflow.keras import backend\nfrom tensorflow.keras import layers\n\nphysical_devices = tf.config.list_physical_devices(\"GPU\")\ntry:\n    tf.config.experimental.set_memory_growth(physical_devices[0], True)\n    tf.config.optimizer.set_jit(True)\n    keras.mixed_precision.set_global_policy(\"mixed_float16\")\nexcept:\n    pass\n\nseed = 1337\ntf.random.set_seed(seed)","metadata":{"id":"idhwz0S2OJMe","execution":{"iopub.status.busy":"2023-11-08T04:12:52.954005Z","iopub.execute_input":"2023-11-08T04:12:52.954356Z","iopub.status.idle":"2023-11-08T04:12:57.744693Z","shell.execute_reply.started":"2023-11-08T04:12:52.954326Z","shell.execute_reply":"2023-11-08T04:12:57.743239Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"**Utils**","metadata":{"id":"9vwk6BIvOJMg"}},{"cell_type":"code","source":"def get_model_weight(model_id):\n    \"\"\"Get the trained weights.\"\"\"\n    if not os.path.exists(\"model.h5\"):\n        model_weight = gdown.download(id=model_id, quiet=False)\n    else:\n        model_weight = \"model.h5\"\n    return model_weight\n\n\ndef get_model_history(history_id):\n    \"\"\"Get the history / log file.\"\"\"\n    if not os.path.exists(\"history.csv\"):\n        history_file = gdown.download(id=history_id, quiet=False)\n    else:\n        history_file = \"history.csv\"\n    return history_file\n\n\ndef make_plot(tfdata, take_batch=1, title=True, figsize=(20, 20)):\n    \"\"\"ref: https://gist.github.com/innat/4dc4080cfdf5cf20ef0fc93d3623ca9b\"\"\"\n\n    font = {\n        \"family\": \"serif\",\n        \"color\": \"darkred\",\n        \"weight\": \"normal\",\n        \"size\": 15,\n    }\n\n    for images, labels in tfdata.take(take_batch):\n        plt.figure(figsize=figsize)\n        xy = int(np.ceil(images.shape[0] * 0.5))\n\n        for i in range(images.shape[0]):\n            plt.subplot(xy, xy, i + 1)\n            plt.imshow(tf.cast(images[i], dtype=tf.uint8))\n            if title:\n                plt.title(tcls_names[tf.argmax(labels[i], axis=-1)], fontdict=font)\n            plt.axis(\"off\")\n\n    plt.tight_layout()\n    plt.show()\n","metadata":{"id":"5XL5UAooOJMh","execution":{"iopub.status.busy":"2023-11-08T04:12:57.747713Z","iopub.execute_input":"2023-11-08T04:12:57.748445Z","iopub.status.idle":"2023-11-08T04:12:57.761976Z","shell.execute_reply.started":"2023-11-08T04:12:57.748411Z","shell.execute_reply":"2023-11-08T04:12:57.760241Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Acquiring Data","metadata":{"id":"RlDh_E8xOJMi"}},{"cell_type":"code","source":"import pathlib\n\ndataset_url = \"https://storage.googleapis.com/download.tensorflow.org/example_images/flower_photos.tgz\"\ndata_dir = keras.utils.get_file('flower_photos', origin=dataset_url, untar=True)\ndata_dir = pathlib.Path(data_dir)\n\nimage_count = len(list(data_dir.glob('*/*.jpg')))\nprint('Total Samples: ', image_count)","metadata":{"id":"guu2tntBOJMj","outputId":"3bf57307-48d1-4ecc-f614-8059d7d15d5b","execution":{"iopub.status.busy":"2023-11-08T04:12:57.763642Z","iopub.execute_input":"2023-11-08T04:12:57.763937Z","iopub.status.idle":"2023-11-08T04:13:01.190152Z","shell.execute_reply.started":"2023-11-08T04:12:57.763913Z","shell.execute_reply":"2023-11-08T04:13:01.188898Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"**Parameter Settings**","metadata":{"id":"kg_sIpHMOJMk"}},{"cell_type":"code","source":"class Parameters:\n    # data level\n    image_size = 384\n    batch_size = 8\n    num_grad_accumulation = 8\n    label_smooth=0.05\n    class_number = 5\n    val_split = 0.2 \n    verbosity = 2\n    autotune = tf.data.AUTOTUNE\n    \n    # hparams\n    epochs = 20\n    lr_sched = 'cosine_restart' # [or, exponential, cosine, linear, constant]\n    lr_base  = 0.016\n    lr_min   = 0\n    lr_decay_epoch  = 2.4\n    lr_warmup_epoch = 5\n    lr_decay_factor = 0.97\n    \n    scaled_lr = lr_base * (batch_size / 256.0)\n    scaled_lr_min = lr_min * (batch_size / 256.0)\n    num_validation_sample = int(image_count * val_split)\n    num_training_sample = image_count - num_validation_sample\n    train_step = int(np.ceil(num_training_sample / float(batch_size)))\n    total_steps = train_step * epochs\n\nparams = Parameters()","metadata":{"id":"bIIjr5xKOJMk","execution":{"iopub.status.busy":"2023-11-08T04:13:01.191238Z","iopub.execute_input":"2023-11-08T04:13:01.191543Z","iopub.status.idle":"2023-11-08T04:13:01.19933Z","shell.execute_reply.started":"2023-11-08T04:13:01.19152Z","shell.execute_reply":"2023-11-08T04:13:01.19792Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Data Loading","metadata":{"id":"7zWGGYIMOJMl"}},{"cell_type":"code","source":"train_set = keras.utils.image_dataset_from_directory(\n    data_dir,\n    validation_split=params.val_split,\n    subset=\"training\",\n    label_mode='categorical',\n    seed=params.image_size,\n    image_size=(params.image_size, params.image_size),\n    batch_size=params.batch_size,\n)\n\nval_set = keras.utils.image_dataset_from_directory(\n    data_dir,\n    validation_split=params.val_split,\n    subset=\"validation\",\n    label_mode='categorical',\n    seed=params.image_size,\n    image_size=(params.image_size, params.image_size),\n    batch_size=params.batch_size,\n)\n\ntcls_names, vcls_names = train_set.class_names , val_set.class_names\ntcls_names, vcls_names ","metadata":{"id":"_NQ2VHOHOJMl","outputId":"fb2bde1d-18fd-4a7f-99f7-790abb7e014f","execution":{"iopub.status.busy":"2023-11-08T04:13:01.201193Z","iopub.execute_input":"2023-11-08T04:13:01.201589Z","iopub.status.idle":"2023-11-08T04:13:02.162307Z","shell.execute_reply.started":"2023-11-08T04:13:01.201561Z","shell.execute_reply":"2023-11-08T04:13:02.160639Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"**Visualize Raw Samples**","metadata":{"id":"ZRyGKHu3OJMm"}},{"cell_type":"code","source":"make_plot(train_set, take_batch=1, title=True) \nmake_plot(val_set, take_batch=1, title=True) ","metadata":{"id":"jvFQOEsTOJMm","execution":{"iopub.status.busy":"2023-11-08T04:13:02.16477Z","iopub.execute_input":"2023-11-08T04:13:02.1652Z","iopub.status.idle":"2023-11-08T04:13:06.342856Z","shell.execute_reply.started":"2023-11-08T04:13:02.165168Z","shell.execute_reply":"2023-11-08T04:13:06.341262Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Advance Image Augmentation\n\nAs mentioned, we will be using **CutMix** and **MixUp**, written in `tf.keras`. And along with it, we will be using two simple `Jax` coded augmentation layers and convert them to work in `tf.data` API. However, currently, there is no such **probability** parameter to control the occurrence of the [keras built-in augmentation](https://keras.io/api/layers/preprocessing_layers/image_augmentation/) layer. So, we will use a wrapper class for the augmentation layers and make the image transformation random in action.","metadata":{"id":"fUD9Jq4uOJMm"}},{"cell_type":"code","source":"class RandomApply(layers.Layer):\n    \"\"\"RandomApply will randomly apply the transformation layer\n    based on the given probability.\n    \n    Ref. https://stackoverflow.com/a/72558994/9215780\n    \"\"\"\n\n    def __init__(self, layer, probability, **kwargs):\n        super().__init__(**kwargs)\n        self.layer = layer\n        self.probability = probability\n\n    def call(self, inputs, training=True):\n        apply_layer = tf.random.uniform([]) < self.probability\n        outputs = tf.cond(\n            pred=tf.logical_and(apply_layer, training),\n            true_fn=lambda: self.layer(inputs),\n            false_fn=lambda: inputs,\n        )\n        return outputs\n\n    def get_config(self):\n        config = super().get_config()\n        config.update(\n            {\n                \"layer\": layers.serialize(self.layer),\n                \"probability\": self.probability,\n            }\n        )\n        return config","metadata":{"id":"aZN7C_eZOJMn","execution":{"iopub.status.busy":"2023-11-08T04:13:06.344418Z","iopub.execute_input":"2023-11-08T04:13:06.344737Z","iopub.status.idle":"2023-11-08T04:13:06.353088Z","shell.execute_reply.started":"2023-11-08T04:13:06.344712Z","shell.execute_reply":"2023-11-08T04:13:06.351756Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## MixUp\n\nImplemented in `tf.keras`.","metadata":{"id":"XMcwYmT2OJMn"}},{"cell_type":"code","source":"class MixUp(layers.Layer):\n    \"\"\"Original implementation: https://github.com/keras-team/keras-cv.\n    The original implementaiton provide more interface to apply mixup on\n    various CV related task, i.e. object detection etc. It also provides\n    many effective validation check.\n\n    Derived and modified for simpler usages: M.Innat.\n    Ref. https://gist.github.com/innat/0ee2b6155d663aac2617fe596e1d8d49\n    \"\"\"\n\n    def __init__(self, alpha=0.2, seed=None, **kwargs):\n        super().__init__(**kwargs)\n        self.alpha = alpha\n        self.seed = seed\n\n    @staticmethod\n    def _sample_from_beta(alpha, beta, shape):\n        sample_alpha = tf.random.gamma(shape, 1.0, beta=alpha)\n        sample_beta = tf.random.gamma(shape, 1.0, beta=beta)\n        return sample_alpha / (sample_alpha + sample_beta)\n\n    def _mixup_samples(self, images):\n        batch_size = tf.shape(images)[0]\n        permutation_order = tf.random.shuffle(tf.range(0, batch_size), seed=self.seed)\n\n        lambda_sample = MixUp._sample_from_beta(self.alpha, self.alpha, (batch_size,))\n        lambda_sample = tf.reshape(lambda_sample, [-1, 1, 1, 1])\n\n        mixup_images = tf.gather(images, permutation_order)\n        images = lambda_sample * images + (1.0 - lambda_sample) * mixup_images\n\n        return images, tf.squeeze(lambda_sample), permutation_order\n\n    def _mixup_labels(self, labels, lambda_sample, permutation_order):\n        labels_for_mixup = tf.gather(labels, permutation_order)\n\n        lambda_sample = tf.reshape(lambda_sample, [-1, 1])\n        labels = lambda_sample * labels + (1.0 - lambda_sample) * labels_for_mixup\n\n        return labels\n\n    def call(self, batch_inputs):\n        bs_images = tf.cast(batch_inputs[0], dtype=tf.float32)  \n        bs_labels = tf.cast(batch_inputs[1], dtype=tf.float32)  \n\n        mixup_images, lambda_sample, permutation_order = self._mixup_samples(bs_images)\n        mixup_labels = self._mixup_labels(bs_labels, lambda_sample, permutation_order)\n\n        return [mixup_images, mixup_labels]\n\n    def get_config(self):\n        config = super().get_config()\n        config.update(\n            {\n                \"alpha\": self.alpha,\n                \"seed\": self.seed,\n            }\n        )\n        return config\n","metadata":{"id":"NzcIrsjOOJMn","execution":{"iopub.status.busy":"2023-11-08T04:13:06.356805Z","iopub.execute_input":"2023-11-08T04:13:06.357367Z","iopub.status.idle":"2023-11-08T04:13:06.37824Z","shell.execute_reply.started":"2023-11-08T04:13:06.357337Z","shell.execute_reply":"2023-11-08T04:13:06.377461Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"temp_ds = train_set.map(\n    lambda x, y: MixUp()([x, y]), num_parallel_calls=params.autotune\n)\nmake_plot(temp_ds, take_batch=1, title=False) ","metadata":{"id":"Yn1tvnmvOJMo","execution":{"iopub.status.busy":"2023-11-08T04:13:06.379372Z","iopub.execute_input":"2023-11-08T04:13:06.379851Z","iopub.status.idle":"2023-11-08T04:13:08.932705Z","shell.execute_reply.started":"2023-11-08T04:13:06.379816Z","shell.execute_reply":"2023-11-08T04:13:08.93191Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## CutMix\n\nImplemented in `tf.keras`.","metadata":{"id":"nVOK8nd-OJMo"}},{"cell_type":"code","source":"class CutMix(layers.Layer):\n    \"\"\"Original implementation: https://github.com/keras-team/keras-cv.\n    The original implementaiton provide more interface to apply mixup on\n    various CV related task, i.e. object detection etc. It also provides\n    many effective validation check.\n    \n    Derived and modified for simpler usages: M.Innat.\n    Ref. https://gist.github.com/innat/0524ee77de17f0601f0dee69aa52c713\n    \"\"\"\n\n    def __init__(self, alpha=1.0, seed=None, **kwargs):\n        super().__init__(**kwargs)\n        self.alpha = alpha\n        self.seed = seed\n\n    @staticmethod\n    def _sample_from_beta(alpha, beta, shape):\n        sample_alpha = tf.random.gamma(shape, 1.0, beta=alpha)\n        sample_beta = tf.random.gamma(shape, 1.0, beta=beta)\n        return sample_alpha / (sample_alpha + sample_beta)\n\n    def _cutmix_labels(self, labels, lambda_sample, permutation_order):\n        cutout_labels = tf.gather(labels, permutation_order)\n\n        lambda_sample = tf.reshape(lambda_sample, [-1, 1])\n        labels = lambda_sample * labels + (1.0 - lambda_sample) * cutout_labels\n        return labels\n\n    def _cutmix_samples(self, images):\n        input_shape = tf.shape(images)\n        batch_size, image_height, image_width = (\n            input_shape[0],\n            input_shape[1],\n            input_shape[2],\n        )\n\n        permutation_order = tf.random.shuffle(tf.range(0, batch_size), seed=self.seed)\n        lambda_sample = CutMix._sample_from_beta(self.alpha, self.alpha, (batch_size,))\n\n        ratio = tf.math.sqrt(1 - lambda_sample)\n\n        cut_height = tf.cast(\n            ratio * tf.cast(image_height, dtype=tf.float32), dtype=tf.int32\n        )\n        cut_width = tf.cast(\n            ratio * tf.cast(image_height, dtype=tf.float32), dtype=tf.int32\n        )\n\n        random_center_height = tf.random.uniform(\n            shape=[batch_size], minval=0, maxval=image_height, dtype=tf.int32\n        )\n        random_center_width = tf.random.uniform(\n            shape=[batch_size], minval=0, maxval=image_width, dtype=tf.int32\n        )\n\n        bounding_box_area = cut_height * cut_width\n        lambda_sample = 1.0 - bounding_box_area / (image_height * image_width)\n        lambda_sample = tf.cast(lambda_sample, dtype=tf.float32)\n\n        images = self.fill_rectangle(\n            images,\n            random_center_width,\n            random_center_height,\n            cut_width,\n            cut_height,\n            tf.gather(images, permutation_order),\n        )\n\n        return images, lambda_sample, permutation_order\n\n    def call(self, batch_inputs, training=None):\n        bs_images = tf.cast(batch_inputs[0], dtype=tf.float32)  \n        bs_labels = tf.cast(batch_inputs[1], dtype=tf.float32)  \n\n        cutmix_images, lambda_sample, permutation_order = self._cutmix_samples(\n            bs_images\n        )\n        cutmix_labels = self._cutmix_labels(bs_labels, lambda_sample, permutation_order)\n\n        return [cutmix_images, cutmix_labels]\n\n    def fill_rectangle(\n        self, images, centers_x, centers_y, widths, heights, fill_values\n    ):\n        images_shape = tf.shape(images)\n        images_height = images_shape[1]\n        images_width = images_shape[2]\n\n        xywh = tf.stack([centers_x, centers_y, widths, heights], axis=1)\n        xywh = tf.cast(xywh, tf.float32)\n        corners = self.convert_format(xywh)\n        mask_shape = (images_width, images_height)\n\n        is_rectangle = self.corners_to_mask(corners, mask_shape)\n        is_rectangle = tf.expand_dims(is_rectangle, -1)\n        images = tf.where(is_rectangle, fill_values, images)\n        return images\n\n    def convert_format(self, boxes):\n        boxes = tf.cast(boxes, dtype=tf.float32)\n        x, y, width, height, rest = tf.split(boxes, [1, 1, 1, 1, -1], axis=-1)\n        results = tf.concat(\n            [\n                x - width / 2.0,\n                y - height / 2.0,\n                x + width / 2.0,\n                y + height / 2.0,\n                rest,\n            ],\n            axis=-1,\n        )\n        return results\n\n    def _axis_mask(self, starts, ends, mask_len):\n        # index range of axis\n        batch_size = tf.shape(starts)[0]\n        axis_indices = tf.range(mask_len, dtype=starts.dtype)\n        axis_indices = tf.expand_dims(axis_indices, 0)\n        axis_indices = tf.tile(axis_indices, [batch_size, 1])\n\n        # mask of index bounds\n        axis_mask = tf.greater_equal(axis_indices, starts) & tf.less(axis_indices, ends)\n        return axis_mask\n\n    def corners_to_mask(self, bounding_boxes, mask_shape):\n        mask_width, mask_height = mask_shape\n        x0, y0, x1, y1 = tf.split(bounding_boxes, [1, 1, 1, 1], axis=-1)\n\n        w_mask = self._axis_mask(x0, x1, mask_width)\n        h_mask = self._axis_mask(y0, y1, mask_height)\n\n        w_mask = tf.expand_dims(w_mask, axis=1)\n        h_mask = tf.expand_dims(h_mask, axis=2)\n        masks = tf.logical_and(w_mask, h_mask)\n        return masks","metadata":{"id":"52VJ5SGwOJMp","execution":{"iopub.status.busy":"2023-11-08T04:13:08.933935Z","iopub.execute_input":"2023-11-08T04:13:08.934343Z","iopub.status.idle":"2023-11-08T04:13:08.952972Z","shell.execute_reply.started":"2023-11-08T04:13:08.934322Z","shell.execute_reply":"2023-11-08T04:13:08.952238Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"temp_ds = train_set.map(\n    lambda x, y: CutMix()([x, y]), num_parallel_calls=params.autotune\n)\nmake_plot(temp_ds, take_batch=1, title=False) ","metadata":{"id":"aPUIsSDdOJMp","execution":{"iopub.status.busy":"2023-11-08T04:13:08.954324Z","iopub.execute_input":"2023-11-08T04:13:08.954595Z","iopub.status.idle":"2023-11-08T04:13:11.520287Z","shell.execute_reply.started":"2023-11-08T04:13:08.954573Z","shell.execute_reply":"2023-11-08T04:13:11.519187Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class RandomMixUpCutMix(layers.Layer):\n    def __init__(self, switch_prob=0.50, **kwargs):\n        super().__init__(**kwargs)\n        self.switch_prob = (\n            switch_prob  # probability of switching between mixup and cutmix\n        )\n        self.mixup = CutMix()\n        self.cutmix = MixUp()\n\n    def call(self, batch_inputs):\n        return tf.cond(\n            tf.less(\n                tf.random.uniform([], minval=0, maxval=1, dtype=tf.float32),\n                tf.cast(self.switch_prob, tf.float32),\n            ),\n            lambda: self.mixup(batch_inputs),\n            lambda: self.cutmix(batch_inputs),\n        )\n\n    def get_config(self):\n        config = {\n            \"switch_prob\": self.switch_prob,\n        }\n        base_config = super().get_config()\n        return dict(list(base_config.items()) + list(config.items()))","metadata":{"id":"aJHW7f0lOJMq","execution":{"iopub.status.busy":"2023-11-08T04:13:11.521914Z","iopub.execute_input":"2023-11-08T04:13:11.522171Z","iopub.status.idle":"2023-11-08T04:13:11.539485Z","shell.execute_reply.started":"2023-11-08T04:13:11.52215Z","shell.execute_reply":"2023-11-08T04:13:11.537304Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Channel Shuffle (jax2tf)\n\nImplemented in `jax`.","metadata":{"id":"_C4uoV9QOJMq"}},{"cell_type":"code","source":"import jax\nfrom jax import jit\nfrom jax import random\nfrom jax import numpy as jnp\nfrom jax.experimental import jax2tf","metadata":{"id":"IcCLA2GeOJMq","execution":{"iopub.status.busy":"2023-11-08T04:13:11.541955Z","iopub.execute_input":"2023-11-08T04:13:11.543518Z","iopub.status.idle":"2023-11-08T04:13:11.707561Z","shell.execute_reply.started":"2023-11-08T04:13:11.543433Z","shell.execute_reply":"2023-11-08T04:13:11.70641Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class RandomChannelShuffle(layers.Layer):\n    \"\"\"Shuffle channels of an input image.\n\n    Ref. https://gist.github.com/innat/35ab35329e2ca890a17556384056334b\n    \"\"\"\n    def __init__(self, groups=3, **kwargs):\n        super().__init__(**kwargs)\n        self.groups = groups\n\n    @partial(jit, static_argnums=0)\n    def _jax_channel_shuffling(self, images):\n        batch_size, height, width, num_channels = images.shape\n\n        if not num_channels % self.groups == 0:\n            raise ValueError(\n                \"The number of input channels should be \"\n                \"divisible by the number of groups.\"\n                f\"Received: channels={num_channels}, groups={self.groups}\"\n            )\n\n        channels_per_group = num_channels // self.groups\n\n        images = images.reshape(-1, height, width, self.groups, channels_per_group)\n        images = images.transpose([3, 1, 2, 4, 0])\n        key = random.PRNGKey(np.random.randint(50))\n        images = random.permutation(key=key, x=images, axis=0)\n        images = images.transpose([4, 1, 2, 3, 0])\n        images = images.reshape(-1, height, width, num_channels)\n        return images\n\n    def call(self, images, training=True):\n        if training:\n            return jax2tf.convert(\n                self._jax_channel_shuffling, polymorphic_shapes=(\"batch, ...\")\n            )(images)\n        else:\n            return images\n    \n    def get_config(self):\n        config = {\n            \"groups\": self.groups,\n        }\n        base_config = super().get_config()\n        return dict(list(base_config.items()) + list(config.items()))","metadata":{"id":"Fg3ijUxOOJMr","execution":{"iopub.status.busy":"2023-11-08T04:13:11.708834Z","iopub.execute_input":"2023-11-08T04:13:11.709134Z","iopub.status.idle":"2023-11-08T04:13:11.72187Z","shell.execute_reply.started":"2023-11-08T04:13:11.709107Z","shell.execute_reply":"2023-11-08T04:13:11.719907Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"temp_ds = train_set.map(\n    lambda x, y: (RandomChannelShuffle()(x), y), num_parallel_calls=params.autotune\n)\nmake_plot(temp_ds, take_batch=1, title=False) ","metadata":{"id":"L2HhSc0QOJMr","execution":{"iopub.status.busy":"2023-11-08T04:13:11.723825Z","iopub.execute_input":"2023-11-08T04:13:11.724352Z","iopub.status.idle":"2023-11-08T04:13:24.44483Z","shell.execute_reply.started":"2023-11-08T04:13:11.724306Z","shell.execute_reply":"2023-11-08T04:13:24.442447Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## GrayScaling (jax2tf)\n\nImplemented in `jax`.","metadata":{"id":"D086fsF8OJMr"}},{"cell_type":"code","source":"class RandomGrayscale(layers.Layer):\n    \"\"\"Grayscale is a preprocessing layer that transforms\n    RGB images to Grayscale images.\n\n    Ref. https://gist.github.com/innat/4e89725ccdcd763e0a6ba19216fd60bf\n    \"\"\"\n\n    def __init__(self, output_channel=1, prob=1, **kwargs):\n        super().__init__(**kwargs)\n        self.output_channel = self._check_input_params(output_channel)\n\n    def _check_input_params(self, output_channels):\n        if output_channels not in [1, 3]:\n            raise ValueError(\n                \"Received invalid argument output_channels. \"\n                f\"output_channels must be in 1 or 3. Got {output_channels}\"\n            )\n        return output_channels\n\n    @partial(jit, static_argnums=0)\n    def _jax_gray_scale(self, images):\n        rgb_weights = jnp.array([0.2989, 0.5870, 0.1140], dtype=images.dtype)\n        grayscale = (rgb_weights * images).sum(axis=-1)\n\n        if self.output_channel == 1:\n            grayscale = jnp.expand_dims(grayscale, axis=-1)\n            return grayscale\n        elif self.output_channel == 3:\n            return jnp.stack([grayscale] * 3, axis=-1)\n        else:\n            raise ValueError(\"Unsupported value for `output_channels`.\")\n\n    def call(self, images, training=True):\n        if training:\n            return jax2tf.convert(\n                self._jax_gray_scale, polymorphic_shapes=(\"batch, ...\")\n            )(images)\n        else:\n            return images\n\n    def get_config(self):\n        config = {\n            \"output_channel\": self.output_channel,\n        }\n        base_config = super().get_config()\n        return dict(list(base_config.items()) + list(config.items()))\n","metadata":{"id":"pZTu5Ru-OJMr","execution":{"iopub.status.busy":"2023-11-08T04:13:24.446609Z","iopub.execute_input":"2023-11-08T04:13:24.446936Z","iopub.status.idle":"2023-11-08T04:13:24.458144Z","shell.execute_reply.started":"2023-11-08T04:13:24.446907Z","shell.execute_reply":"2023-11-08T04:13:24.456965Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"temp_ds = train_set.map(\n    lambda x, y: (RandomGrayscale(output_channel=3)(x), y), num_parallel_calls=params.autotune\n)\nmake_plot(temp_ds, take_batch=1, title=False) ","metadata":{"id":"Rt7Cx8i1OJMs","execution":{"iopub.status.busy":"2023-11-08T04:13:24.459728Z","iopub.execute_input":"2023-11-08T04:13:24.460014Z","iopub.status.idle":"2023-11-08T04:13:26.135929Z","shell.execute_reply.started":"2023-11-08T04:13:24.459992Z","shell.execute_reply":"2023-11-08T04:13:26.135124Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"**Combine All Augmentation**","metadata":{"id":"8CaweldkOJMs"}},{"cell_type":"code","source":"jax_to_keras_augment = keras.Sequential(\n    [\n        RandomApply(RandomGrayscale(output_channel=3), probability=0.2),\n        RandomApply(RandomChannelShuffle(), probability=0.5),\n    ],\n    name=\"jax2keras_augment\",\n)\n\n\ntf_to_keras_augment = keras.Sequential(\n    [\n        RandomApply(layers.RandomFlip(\"horizontal\"), probability=0.5),\n        RandomApply(layers.RandomZoom(0.2, 0.3), probability=0.2),\n        RandomApply(\n            layers.RandomRotation((0.2, 0.3), fill_mode=\"reflect\"), probability=0.8\n        ),\n    ],\n    name=\"tf2keras_augment\",\n)\n","metadata":{"id":"RYTEpVxzOJMs","execution":{"iopub.status.busy":"2023-11-08T04:13:26.137194Z","iopub.execute_input":"2023-11-08T04:13:26.138644Z","iopub.status.idle":"2023-11-08T04:13:26.171843Z","shell.execute_reply.started":"2023-11-08T04:13:26.138579Z","shell.execute_reply":"2023-11-08T04:13:26.171086Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# for train set : augmentation\nkeras_aug = keras.Sequential(\n    [\n        layers.Resizing(height=params.image_size, width=params.image_size),\n        jax_to_keras_augment,\n        tf_to_keras_augment,\n    ],\n    name=\"keras_augment\",\n)\n\ntrain_ds = train_set.shuffle(10 * params.batch_size)\ntrain_ds = train_ds.map(\n    lambda x, y: (keras_aug(x), y), num_parallel_calls=params.autotune\n)\ntrain_ds = train_ds.map(\n    lambda x, y: RandomMixUpCutMix()([x, y]), num_parallel_calls=params.autotune\n)\n","metadata":{"id":"Z0yrtogCOJMs","execution":{"iopub.status.busy":"2023-11-08T04:13:26.173309Z","iopub.execute_input":"2023-11-08T04:13:26.173815Z","iopub.status.idle":"2023-11-08T04:13:31.096726Z","shell.execute_reply.started":"2023-11-08T04:13:26.17379Z","shell.execute_reply":"2023-11-08T04:13:31.095741Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"**Visualize Augmented Training Set**","metadata":{"id":"R4216ZmzOJMt"}},{"cell_type":"code","source":"make_plot(train_ds, take_batch=5, title=False) ","metadata":{"id":"iF15ipIlOJMt","execution":{"iopub.status.busy":"2023-11-08T04:13:31.097931Z","iopub.execute_input":"2023-11-08T04:13:31.09817Z","iopub.status.idle":"2023-11-08T04:13:39.007731Z","shell.execute_reply.started":"2023-11-08T04:13:31.098149Z","shell.execute_reply":"2023-11-08T04:13:39.006137Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"**Visualize Validation Set**","metadata":{"id":"imhUZcfVOJMt"}},{"cell_type":"code","source":"make_plot(val_set, take_batch=1, title=False) ","metadata":{"id":"tk5vrPUqOJMt","execution":{"iopub.status.busy":"2023-11-08T04:13:39.009523Z","iopub.execute_input":"2023-11-08T04:13:39.009848Z","iopub.status.idle":"2023-11-08T04:13:41.517497Z","shell.execute_reply.started":"2023-11-08T04:13:39.009823Z","shell.execute_reply":"2023-11-08T04:13:41.514195Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Overlaps data preprocessing and model execution while training.\n# It often improves latency and throughput, at the cost of using \n# additional memory to store prefetched elements.\ntrain_ds = train_ds.prefetch(buffer_size=params.autotune)\nval_ds = val_set.prefetch(buffer_size=params.autotune)","metadata":{"id":"fArgtRF0OJMt","execution":{"iopub.status.busy":"2023-11-08T04:13:41.523573Z","iopub.execute_input":"2023-11-08T04:13:41.526648Z","iopub.status.idle":"2023-11-08T04:13:41.54433Z","shell.execute_reply.started":"2023-11-08T04:13:41.526569Z","shell.execute_reply":"2023-11-08T04:13:41.541662Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Modeling\n\nWe will use *EfficientNet B0* model and its mid layer, `block6a_expand_activation` as an input of *Swin Transformer Blocks*. The input shape of model *EfficientNet B0* is set `384` and thus we will get output size `24` at layer `block6a_expand_activation`. And this will be the input of **swin-vit**. See the diagram below. Here the **2D CNN output** and **2D Swin Transformer output** are retrieved in order to compute the **GradCAM**.\n\n![](https://user-images.githubusercontent.com/17668390/173227753-36cbf5a5-ee56-4202-bbb2-6020a388b188.png)\n\n<div align=\"center\">\n    <i>Figure: Hybrid-EfficientNet-Swin Transformer.</i>\n</div>","metadata":{"id":"hO8JhuJgOJMu"}},{"cell_type":"code","source":"patch_size      = (2,2)   # 2-by-2 sized patches\ndropout_rate    = 0.5     # Dropout rate\nnum_heads       = 8       # Attention heads\nembed_dim       = 64      # Embedding dimension\nnum_mlp         = 128     # MLP layer size\nqkv_bias        = True    # Convert embedded patches to query, key, and values\nwindow_size     = 2       # Size of attention window\nshift_size      = 1       # Size of shifting window\nimage_dimension = 24      # Initial image size / Input size of the transformer model \n\nnum_patch_x = image_dimension // patch_size[0]\nnum_patch_y = image_dimension // patch_size[1]","metadata":{"id":"PWsNKutOOJMu","execution":{"iopub.status.busy":"2023-11-08T04:13:41.546417Z","iopub.execute_input":"2023-11-08T04:13:41.546966Z","iopub.status.idle":"2023-11-08T04:13:41.558857Z","shell.execute_reply.started":"2023-11-08T04:13:41.546921Z","shell.execute_reply":"2023-11-08T04:13:41.557107Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def window_partition(x, window_size):\n    _, height, width, channels = x.shape\n    patch_num_y = height // window_size\n    patch_num_x = width // window_size\n    x = tf.reshape(\n        x, shape=(-1, patch_num_y, window_size, patch_num_x, window_size, channels)\n    )\n    x = tf.transpose(x, (0, 1, 3, 2, 4, 5))\n    windows = tf.reshape(x, shape=(-1, window_size, window_size, channels))\n    return windows\n\n\ndef window_reverse(windows, window_size, height, width, channels):\n    patch_num_y = height // window_size\n    patch_num_x = width // window_size\n    x = tf.reshape(\n        windows,\n        shape=(-1, patch_num_y, patch_num_x, window_size, window_size, channels),\n    )\n    x = tf.transpose(x, perm=(0, 1, 3, 2, 4, 5))\n    x = tf.reshape(x, shape=(-1, height, width, channels))\n    return x\n\n\nclass DropPath(layers.Layer):\n    def __init__(self, drop_prob=None, **kwargs):\n        super(DropPath, self).__init__(**kwargs)\n        self.drop_prob = drop_prob\n\n    def call(self, inputs, training=None):\n        if self.drop_prob == 0.0 or not training:\n            return inputs\n        else:\n            batch_size = tf.shape(inputs)[0]\n            keep_prob = 1 - self.drop_prob\n            path_mask_shape = (batch_size,) + (1,) * (len(tf.shape(inputs)) - 1)\n            path_mask = tf.floor(\n                backend.random_bernoulli(path_mask_shape, p=keep_prob)\n            )\n            outputs = (\n                tf.math.divide(tf.cast(inputs, dtype=tf.float32), keep_prob) * path_mask\n            )\n            return outputs\n\n    def get_config(self):\n        config = super().get_config()\n        config.update(\n            {\n                \"drop_prob\": self.drop_prob,\n            }\n        )\n        return config","metadata":{"id":"9rwejUxcOJMu","execution":{"iopub.status.busy":"2023-11-08T04:13:41.561798Z","iopub.execute_input":"2023-11-08T04:13:41.562468Z","iopub.status.idle":"2023-11-08T04:13:41.586459Z","shell.execute_reply.started":"2023-11-08T04:13:41.562422Z","shell.execute_reply":"2023-11-08T04:13:41.583424Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class PatchExtract(layers.Layer):\n    def __init__(self, patch_size, **kwargs):\n        super().__init__(**kwargs)\n        self.patch_size_x = patch_size[0]\n        self.patch_size_y = patch_size[0]\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_x, self.patch_size_y, 1),\n            strides=(1, self.patch_size_x, self.patch_size_y, 1),\n            rates=(1, 1, 1, 1),\n            padding=\"VALID\",\n        )\n        patch_dim = patches.shape[-1]\n        patch_num = patches.shape[1]\n        return tf.reshape(patches, (batch_size, patch_num * patch_num, patch_dim))\n\n    def get_config(self):\n        config = super().get_config()\n        config.update(\n            {\n                \"patch_size_y\": self.patch_size_y,\n                \"patch_size_x\": self.patch_size_x,\n            }\n        )\n        return config\n\n\nclass PatchEmbedding(layers.Layer):\n    def __init__(self, num_patch, embed_dim, **kwargs):\n        super().__init__(**kwargs)\n        self.num_patch = num_patch\n        self.proj = layers.Dense(embed_dim)\n        self.pos_embed = layers.Embedding(input_dim=num_patch, output_dim=embed_dim)\n\n    def call(self, patch):\n        pos = tf.range(start=0, limit=self.num_patch, delta=1)\n        return self.proj(patch) + self.pos_embed(pos)\n\n    def get_config(self):\n        config = super().get_config()\n        config.update(\n            {\n                \"num_patch\": self.num_patch,\n            }\n        )\n        return config\n\n\nclass PatchMerging(layers.Layer):\n    def __init__(self, num_patch, embed_dim):\n        super().__init__()\n        self.num_patch = num_patch\n        self.embed_dim = embed_dim\n        self.linear_trans = layers.Dense(2 * embed_dim, use_bias=False)\n\n    def call(self, x):\n        height, width = self.num_patch\n        _, _, C = x.get_shape().as_list()\n        x = tf.reshape(x, shape=(-1, height, width, C))\n        feat_maps = x\n\n        x0 = x[:, 0::2, 0::2, :]\n        x1 = x[:, 1::2, 0::2, :]\n        x2 = x[:, 0::2, 1::2, :]\n        x3 = x[:, 1::2, 1::2, :]\n        x = tf.concat((x0, x1, x2, x3), axis=-1)\n        x = tf.reshape(x, shape=(-1, (height // 2) * (width // 2), 4 * C))\n        return self.linear_trans(x), feat_maps\n\n    def get_config(self):\n        config = super().get_config()\n        config.update({\"num_patch\": self.num_patch, \"embed_dim\": self.embed_dim})\n        return config","metadata":{"id":"3WRPqf-9OJMv","execution":{"iopub.status.busy":"2023-11-08T04:13:41.590457Z","iopub.execute_input":"2023-11-08T04:13:41.590983Z","iopub.status.idle":"2023-11-08T04:13:41.613021Z","shell.execute_reply.started":"2023-11-08T04:13:41.590955Z","shell.execute_reply":"2023-11-08T04:13:41.611494Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class WindowAttention(layers.Layer):\n    def __init__(\n        self,\n        dim,\n        window_size,\n        num_heads,\n        qkv_bias=True,\n        dropout_rate=0.0,\n        return_attention_scores=False,\n        **kwargs\n    ):\n        super().__init__(**kwargs)\n        self.dim = dim\n        self.window_size = window_size\n        self.num_heads = num_heads\n        self.scale = (dim // num_heads) ** -0.5\n        self.return_attention_scores = return_attention_scores\n        self.qkv = layers.Dense(dim * 3, use_bias=qkv_bias)\n        self.dropout = layers.Dropout(dropout_rate)\n        self.proj = layers.Dense(dim)\n\n    def build(self, input_shape):\n        self.relative_position_bias_table = self.add_weight(\n            shape=(\n                (2 * self.window_size[0] - 1) * (2 * self.window_size[1] - 1),\n                self.num_heads,\n            ),\n            initializer=\"zeros\",\n            trainable=True,\n            name=\"relative_position_bias_table\",\n        )\n\n        self.relative_position_index = self.get_relative_position_index(\n            self.window_size[0], self.window_size[1]\n        )\n        super().build(input_shape)\n\n    def get_relative_position_index(self, window_height, window_width):\n        x_x, y_y = tf.meshgrid(range(window_height), range(window_width))\n        coords = tf.stack([y_y, x_x], axis=0)\n        coords_flatten = tf.reshape(coords, [2, -1])\n\n        relative_coords = coords_flatten[:, :, None] - coords_flatten[:, None, :]\n        relative_coords = tf.transpose(relative_coords, perm=[1, 2, 0])\n\n        x_x = (relative_coords[:, :, 0] + window_height - 1) * (2 * window_width - 1)\n        y_y = relative_coords[:, :, 1] + window_width - 1\n        relative_coords = tf.stack([x_x, y_y], axis=-1)\n\n        return tf.reduce_sum(relative_coords, axis=-1)\n\n    def call(self, x, mask=None):\n        _, size, channels = x.shape\n        head_dim = channels // self.num_heads\n        x_qkv = self.qkv(x)\n        x_qkv = tf.reshape(x_qkv, shape=(-1, size, 3, self.num_heads, head_dim))\n        x_qkv = tf.transpose(x_qkv, perm=(2, 0, 3, 1, 4))\n        q, k, v = x_qkv[0], x_qkv[1], x_qkv[2]\n        q = q * self.scale\n        k = tf.transpose(k, perm=(0, 1, 3, 2))\n        attn = q @ k\n\n        relative_position_bias = tf.gather(\n            self.relative_position_bias_table,\n            self.relative_position_index,\n            axis=0,\n        )\n        relative_position_bias = tf.transpose(relative_position_bias, [2, 0, 1])\n        attn = attn + tf.expand_dims(relative_position_bias, axis=0)\n\n        if mask is not None:\n            nW = mask.get_shape()[0]\n            mask_float = tf.cast(\n                tf.expand_dims(tf.expand_dims(mask, axis=1), axis=0), tf.float32\n            )\n            attn = (\n                tf.reshape(attn, shape=(-1, nW, self.num_heads, size, size))\n                + mask_float\n            )\n            attn = tf.reshape(attn, shape=(-1, self.num_heads, size, size))\n            attn = tf.nn.softmax(attn, axis=-1)\n        else:\n            attn = tf.nn.softmax(attn, axis=-1)\n        attn = self.dropout(attn)\n\n        x_qkv = attn @ v\n        x_qkv = tf.transpose(x_qkv, perm=(0, 2, 1, 3))\n        x_qkv = tf.reshape(x_qkv, shape=(-1, size, channels))\n        x_qkv = self.proj(x_qkv)\n        x_qkv = self.dropout(x_qkv)\n\n        if self.return_attention_scores:\n            return x_qkv, attn\n        else:\n            return x_qkv\n\n    def get_config(self):\n        config = super().get_config()\n        config.update(\n            {\n                \"dim\": self.dim,\n                \"window_size\": self.window_size,\n                \"num_heads\": self.num_heads,\n                \"scale\": self.scale,\n            }\n        )\n        return config","metadata":{"id":"ZxjX964JOJMv","execution":{"iopub.status.busy":"2023-11-08T04:13:41.627416Z","iopub.execute_input":"2023-11-08T04:13:41.628086Z","iopub.status.idle":"2023-11-08T04:13:41.661338Z","shell.execute_reply.started":"2023-11-08T04:13:41.628038Z","shell.execute_reply":"2023-11-08T04:13:41.66035Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class SwinTransformer(layers.Layer):\n    def __init__(\n        self, \n        dim,\n        num_patch,\n        num_heads,\n        window_size=7,\n        shift_size=0,\n        num_mlp=1024,\n        qkv_bias=True,\n        dropout_rate=0.0,\n        **kwargs,\n    ):\n        super(SwinTransformer, self).__init__(**kwargs)\n\n        self.dim = dim \n        self.num_patch = num_patch  \n        self.num_heads = num_heads \n        self.window_size = window_size  \n        self.shift_size = shift_size  \n        self.num_mlp = num_mlp  \n\n        self.norm1 = layers.LayerNormalization(epsilon=1e-5)\n        self.attn = WindowAttention(\n            dim,\n            window_size=(self.window_size, self.window_size),\n            num_heads=num_heads,\n            qkv_bias=qkv_bias,\n            dropout_rate=dropout_rate,\n        )\n        self.drop_path = (\n            DropPath(dropout_rate) if dropout_rate > 0.0 else tf.identity\n        )\n        self.norm2 = layers.LayerNormalization(epsilon=1e-5)\n\n        self.mlp = keras.Sequential(\n            [\n                layers.Dense(num_mlp),\n                layers.Activation(keras.activations.gelu),\n                layers.Dropout(dropout_rate),\n                layers.Dense(dim),\n                layers.Dropout(dropout_rate),\n            ]\n        )\n\n        if min(self.num_patch) < self.window_size:\n            self.shift_size = 0\n            self.window_size = min(self.num_patch)\n\n    def build(self, input_shape):\n        if self.shift_size == 0:\n            self.attn_mask = None\n        else:\n            height, width = self.num_patch\n            h_slices = (\n                slice(0, -self.window_size),\n                slice(-self.window_size, -self.shift_size),\n                slice(-self.shift_size, None),\n            )\n            w_slices = (\n                slice(0, -self.window_size),\n                slice(-self.window_size, -self.shift_size),\n                slice(-self.shift_size, None),\n            )\n            mask_array = jnp.zeros((1, height, width, 1))\n            count = 0\n            for h in h_slices:\n                for w in w_slices:\n                    mask_array[:, h, w, :] = count\n                    count += 1\n            mask_array = tf.convert_to_tensor(mask_array)\n\n            # mask array to windows\n            mask_windows = window_partition(mask_array, self.window_size)\n            mask_windows = tf.reshape(\n                mask_windows, shape=[-1, self.window_size * self.window_size]\n            )\n            attn_mask = tf.expand_dims(mask_windows, axis=1) - tf.expand_dims(\n                mask_windows, axis=2\n            )\n            attn_mask = tf.where(attn_mask != 0, -100.0, attn_mask)\n            attn_mask = tf.where(attn_mask == 0, 0.0, attn_mask)\n            self.attn_mask = tf.Variable(initial_value=attn_mask, trainable=False)\n\n    def call(self, x):\n        height, width = self.num_patch\n        _, num_patches_before, channels = x.shape\n        x_skip = x\n        x = self.norm1(x)\n        x = tf.reshape(x, shape=(-1, height, width, channels))\n        if self.shift_size > 0:\n            shifted_x = tf.roll(\n                x, shift=[-self.shift_size, -self.shift_size], axis=[1, 2]\n            )\n        else:\n            shifted_x = x\n\n        x_windows = window_partition(shifted_x, self.window_size)\n        x_windows = tf.reshape(\n            x_windows, shape=(-1, self.window_size * self.window_size, channels)\n        )\n        attn_windows = self.attn(x_windows, mask=self.attn_mask)\n\n        attn_windows = tf.reshape(\n            attn_windows, shape=(-1, self.window_size, self.window_size, channels)\n        )\n        shifted_x = window_reverse(\n            attn_windows, self.window_size, height, width, channels\n        )\n        if self.shift_size > 0:\n            x = tf.roll(\n                shifted_x, shift=[self.shift_size, self.shift_size], axis=[1, 2]\n            )\n        else:\n            x = shifted_x\n\n        x = tf.reshape(x, shape=(-1, height * width, channels))\n        x = self.drop_path(x)\n        x = tf.cast(x_skip, dtype=tf.float32) + tf.cast(x, dtype=tf.float32)\n        x_skip = x\n        x = self.norm2(x)\n        x = self.mlp(x)\n        x = self.drop_path(x)\n        x = tf.cast(x_skip, dtype=tf.float32) + tf.cast(x, dtype=tf.float32)\n        return x","metadata":{"id":"HYsF6ndpOJMv","execution":{"iopub.status.busy":"2023-11-08T04:13:41.663711Z","iopub.execute_input":"2023-11-08T04:13:41.665188Z","iopub.status.idle":"2023-11-08T04:13:41.701154Z","shell.execute_reply.started":"2023-11-08T04:13:41.665128Z","shell.execute_reply":"2023-11-08T04:13:41.698337Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Hybrid-EfficientNet-Swin-Transformer","metadata":{"id":"3Nwb669SOJMw"}},{"cell_type":"code","source":"# base cnn models\nbase = keras.applications.EfficientNetB0(\n    include_top=False,\n    weights='imagenet',\n    input_tensor=keras.Input((params.image_size, params.image_size, 3)),\n)\n","metadata":{"id":"2WLyUIlHOJMw","outputId":"da1ebece-9467-4d7c-a149-1e84e70e156f","execution":{"iopub.status.busy":"2023-11-08T04:13:41.704593Z","iopub.execute_input":"2023-11-08T04:13:41.70544Z","iopub.status.idle":"2023-11-08T04:13:44.697043Z","shell.execute_reply.started":"2023-11-08T04:13:41.705393Z","shell.execute_reply":"2023-11-08T04:13:44.695101Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class HybridModel(keras.Model):\n    def __init__(self, model_name, **kwargs):\n        super().__init__(name=model_name, **kwargs)\n\n        # base model with compatible output which will be an input of transformer model\n        self.multi_output_cnn = keras.Model(\n            [base.inputs],\n            [base.get_layer(\"block6a_expand_activation\").output, base.output],\n            name=\"efficientnet\",\n        )\n\n        # base model's (cnn model) head\n        self.conv_head = keras.Sequential(\n            [\n                layers.GlobalAveragePooling2D(),\n                layers.AlphaDropout(0.5),\n            ],\n            name=\"conv_head\",\n        )\n\n        # stuff of swin transformers\n        self.patch_extract = PatchExtract(patch_size)\n        self.patch_embedds = PatchEmbedding(num_patch_x * num_patch_y, embed_dim)\n        self.patch_merging = PatchMerging(\n            (num_patch_x, num_patch_y), embed_dim=embed_dim\n        )\n\n        # swin blocks containers\n        self.swin_sequences = keras.Sequential(name=\"swin_blocks\")\n        for i in range(shift_size):\n            self.swin_sequences.add(\n                SwinTransformer(\n                    dim=embed_dim,\n                    num_patch=(num_patch_x, num_patch_y),\n                    num_heads=num_heads,\n                    window_size=window_size,\n                    shift_size=i,\n                    num_mlp=num_mlp,\n                    qkv_bias=qkv_bias,\n                    dropout_rate=dropout_rate,\n                )\n            )\n\n        # swin block's head\n        self.swin_head = keras.Sequential(\n            [\n                layers.GlobalAveragePooling1D(),\n                layers.AlphaDropout(0.5),\n                layers.BatchNormalization(),\n            ],\n            name=\"swin_head\",\n        )\n\n        # classifier\n        self.classifier = layers.Dense(\n            params.class_number, activation=None, dtype=\"float32\"\n        )\n\n        # build the graph\n        self.build_graph()\n\n    def forward_cnn(self, inputs):\n        # CNN model.\n        return self.multi_output_cnn(inputs)\n\n    def forward_transformer(self, inputs):\n        # Transformer model.\n        x = self.patch_extract(inputs)\n        x = self.patch_embedds(x)\n        x = self.swin_sequences(tf.cast(x, dtype=tf.float32))\n        x, swin_gcam_top = self.patch_merging(x)\n        return x, swin_gcam_top\n\n    def call(self, inputs, training=None, **kwargs):\n        cnn_mid_layer, cnn_gcam_top = self.forward_cnn(inputs)\n        transformer_output, transformer_gcam_top = self.forward_transformer(\n            cnn_mid_layer\n        )\n\n        transformer_output = self.swin_head(transformer_output)\n        cnn_output = self.conv_head(cnn_gcam_top)\n        logits = self.classifier(tf.concat([transformer_output, cnn_output], axis=-1))\n\n        if training:\n            return logits\n        else:\n            return logits, cnn_gcam_top, transformer_gcam_top\n\n    def build_graph(self):\n        x = keras.Input(shape=(params.image_size, params.image_size, 3))\n        return keras.Model(inputs=[x], outputs=self.call(x))\n","metadata":{"id":"wYv-BqNZOJMw","execution":{"iopub.status.busy":"2023-11-08T04:13:44.699821Z","iopub.execute_input":"2023-11-08T04:13:44.700294Z","iopub.status.idle":"2023-11-08T04:13:44.724005Z","shell.execute_reply.started":"2023-11-08T04:13:44.700253Z","shell.execute_reply":"2023-11-08T04:13:44.722144Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Implement Gradient Accumulation\n\nWe also like to implement gradient accumulation technique in our model building pipelines. Usually transformer based models are computationally expensive and thus constrain the batch size limit. To overcome, we like to split up the batch into smaller mini-batches which are run sequentially, while accumulating their results. Because gradient accumulation technique calculates the loss and gradients after each mini-batch, but instead of updating the model parameters, it waits and accumulates the gradients over consecutive batches, so it can overcoming memory constraints, i.e using less memory to training the model like it using large batch size. So, if we run a gradient accumulation with steps of **8** and batch size of **8** images, it serves almost the same purpose of running with a batch size of **64** images.\n\n\n<img src=\"https://miro.medium.com/max/1050/1*rJIH9gPhctTLCk5G5iQ_oA.png\" width=\"500\" height=\"500\" />\n\n[source.](https://towardsdatascience.com/what-is-gradient-accumulation-in-deep-learning-ec034122cfa)","metadata":{"id":"pSCpdkbROJMx"}},{"cell_type":"code","source":"class GradientAccumulation(HybridModel):\n    \"\"\"Subclassing the mdoel class to override the train step to \n    implemnet gradient accumulation.\n    \n    Ref: https://gist.github.com/innat/ba6740293e7b7b227829790686f2119c\n    \"\"\"\n\n    def __init__(self, n_gradients, **kwargs):\n        super().__init__(**kwargs)\n        self.n_gradients = tf.constant(n_gradients, dtype=tf.int32)\n        self.n_acum_step = tf.Variable(0, dtype=tf.int32, trainable=False)\n        self.gradient_accumulation = [\n            tf.Variable(tf.zeros_like(v, dtype=tf.float32), trainable=False)\n            for v in self.trainable_variables\n        ]\n\n    def train_step(self, data):\n        # track accumulation step update\n        self.n_acum_step.assign_add(1)\n\n        # Unpack the data. Its structure depends on your model and\n        # on what you pass to `fit()`.\n        x, y = data\n\n        with tf.GradientTape() as tape:\n            y_pred = self(x, training=True)  # Forward pass\n            loss = self.compiled_loss(y, y_pred, regularization_losses=self.losses)\n\n        # Calculate batch gradients\n        gradients = tape.gradient(loss, self.trainable_variables)\n\n        # Accumulate batch gradients\n        for i in range(len(self.gradient_accumulation)):\n            self.gradient_accumulation[i].assign_add(gradients[i])\n\n        # If n_acum_step reach the n_gradients then we apply accumulated gradients to -\n        # update the variables otherwise do nothing\n        tf.cond(\n            tf.equal(self.n_acum_step, self.n_gradients),\n            self.apply_accu_gradients,\n            lambda: None,\n        )\n\n        # Return a dict mapping metric names to current value.\n        # Note that it will include the loss (tracked in self.metrics).\n        self.compiled_metrics.update_state(y, y_pred)\n        return {m.name: m.result() for m in self.metrics}\n\n    def apply_accu_gradients(self):\n        # Update weights\n        self.optimizer.apply_gradients(\n            zip(self.gradient_accumulation, self.trainable_variables)\n        )\n\n        # reset accumulation step\n        self.n_acum_step.assign(0)\n        for i in range(len(self.gradient_accumulation)):\n            self.gradient_accumulation[i].assign(\n                tf.zeros_like(self.trainable_variables[i], dtype=tf.float32)\n            )\n\n    def test_step(self, data):\n        # Unpack the data\n        x, y = data\n\n        # Compute predictions\n        y_pred, base_gcam_top, swin_gcam_top = self(x, training=False)\n\n        # Updates the metrics tracking the loss\n        self.compiled_loss(y, y_pred, regularization_losses=self.losses)\n\n        # Update the metrics.\n        self.compiled_metrics.update_state(y, y_pred)\n\n        # Return a dict mapping metric names to current value.\n        # Note that it will include the loss (tracked in self.metrics).\n        return {m.name: m.result() for m in self.metrics}\n    ","metadata":{"id":"4HyN9h36OJMx","execution":{"iopub.status.busy":"2023-11-08T04:13:44.728068Z","iopub.execute_input":"2023-11-08T04:13:44.729552Z","iopub.status.idle":"2023-11-08T04:13:44.797327Z","shell.execute_reply.started":"2023-11-08T04:13:44.729497Z","shell.execute_reply":"2023-11-08T04:13:44.795526Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def get_model(plot_summary=False, plot_graph=False):\n    keras.backend.clear_session()\n    model = GradientAccumulation(\n        n_gradients=params.num_grad_accumulation, model_name=\"HybridModel\"\n    )\n\n    if plot_summary:\n        display(\n            model.build_graph().summary()\n        )\n\n    if plot_graph:\n        display(\n            keras.utils.plot_model(\n                model.build_graph(),\n                show_shapes=True,\n                show_layer_names=True,\n                expand_nested=False,\n            )\n        )\n\n    # compile\n    model.compile(\n        loss=losses.CategoricalCrossentropy(\n            label_smoothing=params.label_smooth, from_logits=True\n        ),\n        optimizer=optimizers.Adam(learning_rate, amsgrad=True),\n        metrics=[\"accuracy\"],\n    )\n    return model","metadata":{"id":"X7JhoFJUOJMx","execution":{"iopub.status.busy":"2023-11-08T04:13:44.799895Z","iopub.execute_input":"2023-11-08T04:13:44.800543Z","iopub.status.idle":"2023-11-08T04:13:44.818057Z","shell.execute_reply.started":"2023-11-08T04:13:44.800503Z","shell.execute_reply":"2023-11-08T04:13:44.81645Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"**Pre Setting for Training**","metadata":{"id":"KdiBwyUJOJMx"}},{"cell_type":"code","source":"from tensorflow.keras import losses\nfrom tensorflow.keras import metrics\nfrom tensorflow.keras import callbacks\nfrom tensorflow.keras import optimizers\n\nckp = callbacks.ModelCheckpoint(\n    \"model.h5\",\n    monitor=\"val_accuracy\",\n    verbose=1,\n    save_best_only=True,\n    save_weights_only=True,\n    mode=\"max\",\n)\nlog = callbacks.CSVLogger(\"history.csv\", separator=\",\", append=False)","metadata":{"id":"EhziI0YvOJMy","execution":{"iopub.status.busy":"2023-11-08T04:13:44.819751Z","iopub.execute_input":"2023-11-08T04:13:44.820194Z","iopub.status.idle":"2023-11-08T04:13:44.833526Z","shell.execute_reply.started":"2023-11-08T04:13:44.820156Z","shell.execute_reply":"2023-11-08T04:13:44.831627Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class WarmupLearningRateSchedule(optimizers.schedules.LearningRateSchedule):\n    \"\"\"WarmupLearningRateSchedule a variety of learning rate\n    decay schedules with warm up.\n    \n    Ref. https://gist.github.com/innat/69e8f3500c2418c69b150a0a651f31dc\n    \"\"\"\n\n    def __init__(\n        self,\n        initial_lr,\n        steps_per_epoch=None,\n        lr_decay_type=\"exponential\",\n        decay_factor=0.97,\n        decay_epochs=2.4,\n        total_steps=None,\n        warmup_epochs=5,\n        minimal_lr=0, \n        **kwargs\n    ):\n        super().__init__(**kwargs)\n        self.initial_lr = initial_lr\n        self.steps_per_epoch = steps_per_epoch\n        self.lr_decay_type = lr_decay_type\n        self.decay_factor = decay_factor\n        self.decay_epochs = decay_epochs\n        self.total_steps = total_steps\n        self.warmup_epochs = warmup_epochs\n        self.minimal_lr = minimal_lr\n\n    def __call__(self, step):\n        if self.lr_decay_type == \"exponential\":\n            assert self.steps_per_epoch is not None\n            decay_steps = self.steps_per_epoch * self.decay_epochs\n            lr = schedules.ExponentialDecay(\n                self.initial_lr, decay_steps, self.decay_factor, staircase=True\n            )(step)\n            \n        elif self.lr_decay_type == \"cosine\":\n            assert self.total_steps is not None\n            lr = (\n                0.5\n                * self.initial_lr\n                * (1 + tf.cos(np.pi * tf.cast(step, tf.float32) / self.total_steps))\n            )\n\n        elif self.lr_decay_type == \"linear\":\n            assert self.total_steps is not None\n            lr = (1.0 - tf.cast(step, tf.float32) / self.total_steps) * self.initial_lr\n\n        elif self.lr_decay_type == \"constant\":\n            lr = self.initial_lr\n\n        elif self.lr_decay_type == \"cosine_restart\":\n            decay_steps = self.steps_per_epoch * self.decay_epochs\n            lr = tf.keras.experimental.CosineDecayRestarts(\n                self.initial_lr, decay_steps\n            )(step)\n        else:\n            assert False, \"Unknown lr_decay_type : %s\" % self.lr_decay_type\n\n        if self.minimal_lr:\n            lr = tf.math.maximum(lr, self.minimal_lr)\n\n        if self.warmup_epochs:\n            warmup_steps = int(self.warmup_epochs * self.steps_per_epoch)\n            warmup_lr = (\n                self.initial_lr\n                * tf.cast(step, tf.float32)\n                / tf.cast(warmup_steps, tf.float32)\n            )\n            lr = tf.cond(step < warmup_steps, lambda: warmup_lr, lambda: lr)\n\n        return lr\n\n    def get_config(self):\n        return {\n            \"initial_lr\": self.initial_lr,\n            \"steps_per_epoch\": self.steps_per_epoch,\n            \"lr_decay_type\": self.lr_decay_type,\n            \"decay_factor\": self.decay_factor,\n            \"decay_epochs\": self.decay_epochs,\n            \"total_steps\": self.total_steps,\n            \"warmup_epochs\": self.warmup_epochs,\n            \"minimal_lr\": self.minimal_lr,\n        }\n","metadata":{"id":"2gPbynD7OJMy","execution":{"iopub.status.busy":"2023-11-08T04:13:44.835672Z","iopub.execute_input":"2023-11-08T04:13:44.838195Z","iopub.status.idle":"2023-11-08T04:13:44.857767Z","shell.execute_reply.started":"2023-11-08T04:13:44.838132Z","shell.execute_reply":"2023-11-08T04:13:44.855844Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"learning_rate = WarmupLearningRateSchedule(\n    params.scaled_lr,\n    steps_per_epoch=params.train_step,\n    decay_epochs=params.lr_decay_epoch,\n    warmup_epochs=params.lr_warmup_epoch,\n    decay_factor=params.lr_decay_factor,\n    lr_decay_type=params.lr_sched,\n    total_steps=params.total_steps,\n    minimal_lr=params.scaled_lr_min,\n)\n\nrng = [i for i in range(params.total_steps)]\nlr_y = [learning_rate(x) for x in rng]\nplt.figure(figsize=(10, 4))\nplt.plot(rng, lr_y)\nplt.xlabel(\"Iteration\", size=14)\nplt.ylabel(\"Learning Rate\", size=14)","metadata":{"id":"i3ahmKWmOJM2","outputId":"4c03713a-7d97-4e82-c136-e6dd883671a0","execution":{"iopub.status.busy":"2023-11-08T04:13:44.860143Z","iopub.execute_input":"2023-11-08T04:13:44.86105Z","iopub.status.idle":"2023-11-08T04:13:53.709938Z","shell.execute_reply.started":"2023-11-08T04:13:44.861013Z","shell.execute_reply":"2023-11-08T04:13:53.708691Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# training if GPU is available\nif physical_devices:\n    model = get_model(plot_summary=True, plot_graph=False)\n    history = model.fit(\n        train_ds,\n        epochs=params.epochs,\n        callbacks=[ckp, log],\n        validation_data=val_ds,\n        verbose=params.verbosity,\n    ).history\n\n    # Check if the model file already exists, and optionally overwrite it\n    model_file_path = \"./model.h5\"\n    model.save_weights(model_file_path)\n    display(pd.DataFrame(history).tail(5))\nelse: # inference instead\n    keras.mixed_precision.set_global_policy(\"float32\")\n    model = get_model(plot_summary=True, plot_graph=False)\n    model(tf.ones((1, params.image_size, params.image_size, 3)))[0].shape\n\n    # get trained weight and history file\n    weight = get_model_weight(model_id=\"1y6tseN0194T6d-4iIh5wo7RL9ttQERe0\")\n    model.load_weights(weight)\n\n    history_csv = get_model_history(history_id=\"1J6QgHUqtYL0mIWC2h0K6HeIHcollfRUe\")\n    history = pd.read_csv(history_csv)\n    display(history.tail())\n","metadata":{"id":"xglsLP3IOJM3","outputId":"f58d0491-04fa-41b6-91e5-8e258c943f77","execution":{"iopub.status.busy":"2023-11-08T04:13:53.713294Z","iopub.execute_input":"2023-11-08T04:13:53.713688Z","iopub.status.idle":"2023-11-08T04:14:05.389106Z","shell.execute_reply.started":"2023-11-08T04:13:53.713662Z","shell.execute_reply":"2023-11-08T04:14:05.387793Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"**Learning Curve**","metadata":{"id":"WWl8k0N6OJM3"}},{"cell_type":"code","source":"# Plotting\nplt.figure(figsize=(20, 10))\nplt.plot(range(len(history[\"loss\"])), history[\"loss\"], \"-o\", label=\"train_loss\")\nplt.plot(range(len(history[\"loss\"])), history[\"val_loss\"], \"-o\", label=\"val_loss\")\nplt.plot(range(len(history[\"loss\"])), history[\"accuracy\"], \"-o\", label=\"train_accuracy\")\nplt.plot(\n    range(len(history[\"loss\"])), history[\"val_accuracy\"], \"-o\", label=\"val_accuracy\"\n)\nplt.title(\"Training Loss and Accuracy\", fontdict={'fontsize':20})\nplt.xlabel(f\"Epoch {len(history)}\", fontsize=20)\nplt.ylabel(\"Loss/Accuracy\", fontsize=20)\nplt.legend(loc=\"best\", prop={\"size\": 20})\nplt.tight_layout()\nplt.show()","metadata":{"id":"ZehWYOmVOJM3","outputId":"7c70ba85-1998-4d85-9ea0-d38f9571dfd5","execution":{"iopub.status.busy":"2023-11-08T04:14:05.390867Z","iopub.execute_input":"2023-11-08T04:14:05.391237Z","iopub.status.idle":"2023-11-08T04:14:05.866149Z","shell.execute_reply.started":"2023-11-08T04:14:05.391185Z","shell.execute_reply":"2023-11-08T04:14:05.863694Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Grad-CAM : Hybrid-EfficientNet-Swin-Transformer","metadata":{"id":"bqGru9aeOJM3"}},{"cell_type":"code","source":"def plot_stuff(inputs, features_a, features_b):\n    plt.figure(figsize=(25, 25))\n    \n    plt.subplot(1, 3, 1)\n    plt.axis('off')\n    plt.imshow(tf.squeeze(inputs/255, axis=0))\n    plt.title('Input')\n    \n    plt.subplot(1, 3, 2)\n    plt.axis('off')\n    plt.imshow(features_a)\n    plt.title('CNN')\n    \n    plt.subplot(1, 3, 3)\n    plt.axis('off')\n    plt.imshow(features_b)\n    plt.title('Hybrid-CNN-Transformer')\n    plt.show()\n\n# ref: https://keras.io/examples/vision/grad_cam/\ndef get_img_array(img):\n    array = keras.utils.img_to_array(img)\n    array = np.expand_dims(array, axis=0)\n    return array\n\n# ref: https://keras.io/examples/vision/grad_cam/\ndef make_gradcam_heatmap(img_array, grad_model, pred_index=None):\n    with tf.GradientTape(persistent=True) as tape:\n        preds, base_top, swin_top = grad_model(img_array)\n        if pred_index is None:\n            pred_index = tf.argmax(preds[0])\n        class_channel = preds[:, pred_index]\n\n    grads = tape.gradient(class_channel, base_top)\n    pooled_grads = tf.reduce_mean(grads, axis=(0, 1, 2))\n    base_top = base_top[0]\n    heatmap_a = base_top @ pooled_grads[..., tf.newaxis]\n    heatmap_a = tf.squeeze(heatmap_a)\n    heatmap_a = tf.maximum(heatmap_a, 0) / tf.math.reduce_max(heatmap_a)\n    heatmap_a = heatmap_a.numpy()\n\n    grads = tape.gradient(class_channel, swin_top)\n    pooled_grads = tf.reduce_mean(grads, axis=(0, 1, 2))\n    swin_top = swin_top[0]\n    heatmap_b = swin_top @ pooled_grads[..., tf.newaxis]\n    heatmap_b = tf.squeeze(heatmap_b)\n    heatmap_b = tf.maximum(heatmap_b, 0) / tf.math.reduce_max(heatmap_b)\n    heatmap_b = heatmap_b.numpy()\n    return heatmap_a, heatmap_b","metadata":{"id":"hC7iiJheOJM3","execution":{"iopub.status.busy":"2023-11-08T04:14:05.868322Z","iopub.execute_input":"2023-11-08T04:14:05.868747Z","iopub.status.idle":"2023-11-08T04:14:05.880927Z","shell.execute_reply.started":"2023-11-08T04:14:05.868711Z","shell.execute_reply":"2023-11-08T04:14:05.879953Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Get val images\nimg_arrays = next(iter(val_ds))[0]\nprint(img_arrays.shape)\n\n# plot utils\nfor img_array in img_arrays[:3]:\n    # Generate class activation heatmap\n    img_array = get_img_array(img_array)\n    cnn_heatmap, swin_heatmap = make_gradcam_heatmap(img_array, model)\n    print(cnn_heatmap.shape, cnn_heatmap.max(), cnn_heatmap.min())\n    print(swin_heatmap.shape, swin_heatmap.max(), swin_heatmap.min())\n\n    # Display heatmap\n    plot_stuff(img_array, cnn_heatmap, swin_heatmap)","metadata":{"id":"1N9WvaTBOJM4","outputId":"fc455640-deaa-454c-cef7-6ce26aa01c26","execution":{"iopub.status.busy":"2023-11-08T04:14:05.882363Z","iopub.execute_input":"2023-11-08T04:14:05.882922Z","iopub.status.idle":"2023-11-08T04:14:12.757258Z","shell.execute_reply.started":"2023-11-08T04:14:05.882894Z","shell.execute_reply":"2023-11-08T04:14:12.756281Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# ref: https://keras.io/examples/vision/grad_cam/\ndef save_and_display_gradcam(\n    img,\n    heatmap,\n    target=None,\n    pred=None,\n    cam_path=\"cam.jpg\",\n    cmap=\"jet\",  # inferno, viridis\n    alpha=0.6,\n    plot=None,\n):\n    # Rescale heatmap to a range 0-255\n    heatmap = np.uint8(255 * heatmap)\n\n    # Use jet colormap to colorize heatmap\n    jet = cm.get_cmap(cmap)\n\n    # Use RGB values of the colormap\n    jet_colors = jet(np.arange(256))[:, :3]\n    jet_heatmap = jet_colors[heatmap]\n\n    # Create an image with RGB colorized heatmap\n    jet_heatmap = keras.utils.array_to_img(jet_heatmap)\n    jet_heatmap = jet_heatmap.resize((img.shape[0], img.shape[1]))\n    jet_heatmap = keras.utils.img_to_array(jet_heatmap)\n\n    # Superimpose the heatmap on original image\n    superimposed_img = img + jet_heatmap * alpha\n    superimposed_img = keras.utils.array_to_img(superimposed_img)\n    return superimposed_img","metadata":{"id":"bJ4QiCsaOJM4","execution":{"iopub.status.busy":"2023-11-08T04:14:12.758661Z","iopub.execute_input":"2023-11-08T04:14:12.759826Z","iopub.status.idle":"2023-11-08T04:14:12.769337Z","shell.execute_reply.started":"2023-11-08T04:14:12.759792Z","shell.execute_reply":"2023-11-08T04:14:12.766736Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"samples, labels = next(iter(val_ds.shuffle(params.batch_size)))\n\nfor sample, label in zip(samples, labels):\n    # preparing\n    img_array = sample[tf.newaxis, ...]\n\n    # get heatmaps\n    heatmap_a, heatmap_b = make_gradcam_heatmap(img_array, model)\n\n    # overaly heatmap and input sample\n    overaly_a = save_and_display_gradcam(sample, heatmap_a)\n    overlay_b = save_and_display_gradcam(sample, heatmap_b)\n\n    # ploting stuff\n    plot_stuff(img_array, overaly_a, overlay_b)","metadata":{"id":"OBPjGQ8VOJM4","outputId":"07c7951e-e604-441b-9703-99a6622452da","execution":{"iopub.status.busy":"2023-11-08T04:14:12.77084Z","iopub.execute_input":"2023-11-08T04:14:12.771354Z","iopub.status.idle":"2023-11-08T04:14:35.292708Z","shell.execute_reply.started":"2023-11-08T04:14:12.77132Z","shell.execute_reply":"2023-11-08T04:14:35.291525Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Model Saving and Reloading\n\nPreviously, we have saved the model's weights using `callbacks.ModelCheckpoint` API. But we can also save the entire model in `TensorFlow SavedModel` format (recommended format option). A `SavedModel` contains a complete TensorFlow program, including trained parameters (i.e, [tf.Variables](https://www.tensorflow.org/api_docs/python/tf/Variable)) and computation. It does not require the original model building code to run, which makes it useful for sharing or deploying with [TFLite](https://www.tensorflow.org/lite), [TensorFlow.js](https://www.tensorflow.org/js/), [TensorFlow Serving](https://www.tensorflow.org/tfx/tutorials/serving/rest_simple), or [TensorFlow Hub](https://www.tensorflow.org/hub).","metadata":{"id":"WJH88swMOJM4"}},{"cell_type":"code","source":"# Calling `save('my_model')` creates a SavedModel folder `my_model`.\nmodel.save(\"saved_model\")\n\n\n# It can be used to reconstruct the model identically.\nreconstructed_model = keras.models.load_model(\n    \"saved_model\",\n    custom_objects={\"WarmupLearningRateSchedule\": WarmupLearningRateSchedule},\n)\n","metadata":{"id":"aZU6tLweOJM4","outputId":"1a20cf29-eda1-42ee-b6f8-da7cd1383d49","execution":{"iopub.status.busy":"2023-11-08T04:14:35.29417Z","iopub.execute_input":"2023-11-08T04:14:35.29453Z","iopub.status.idle":"2023-11-08T04:15:52.602526Z","shell.execute_reply.started":"2023-11-08T04:14:35.294502Z","shell.execute_reply":"2023-11-08T04:15:52.601308Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Let's check: weight matching\nassert len(model.weights) == len(reconstructed_model.weights)\nfor a, b in zip(model.weights, reconstructed_model.weights):\n    np.testing.assert_allclose(a.numpy(), b.numpy())\n\n    \n# Let's check: inference matching\ntest_input = tf.random.normal(\n    [1, params.image_size, params.image_size, 3], 0, 1, tf.float32\n)\ntf.nest.map_structure(\n    np.testing.assert_allclose,\n    model.predict(test_input),\n    reconstructed_model.predict(test_input),\n)\n","metadata":{"id":"nCBH_Q68OJM5","outputId":"0c4c233c-545f-41bc-adee-4cc8d99702e5","execution":{"iopub.status.busy":"2023-11-08T04:15:52.623501Z","iopub.execute_input":"2023-11-08T04:15:52.624072Z","iopub.status.idle":"2023-11-08T04:15:56.759992Z","shell.execute_reply.started":"2023-11-08T04:15:52.624031Z","shell.execute_reply":"2023-11-08T04:15:56.758972Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# TensorFlow Lite Conversion\n\nThis **HybridSwinTransformer** model can be converted into [TensorFlow Lite](https://www.tensorflow.org/lite/guide) format for mobile and edge devices. Below is the conversion code. Note that, the `TFLite` can be used to minimize the complexity of optimizing inference. You can read more details about it from [here](https://www.tensorflow.org/lite/performance/model_optimization).","metadata":{"id":"Pn7xKprxOJM5"}},{"cell_type":"code","source":"from tensorflow import lite\n\n# wrap keras model, optimize, enable tf operation\n# and convert.\nconverter = lite.TFLiteConverter.from_keras_model(model)\nconverter.optimizations = [tf.lite.Optimize.DEFAULT]\nconverter.target_spec.supported_ops = [\n    lite.OpsSet.TFLITE_BUILTINS,\n    lite.OpsSet.SELECT_TF_OPS,\n]\nconverter.allow_custom_ops = True\ntflite_model = converter.convert()\n\n# Save the model.\nwith open(\"model.tflite\", \"wb\") as write_tflite:\n    write_tflite.write(tflite_model)","metadata":{"id":"TgVNUjaQOJM5","outputId":"155eabfb-4611-4166-8cae-1eaf56d43983","execution":{"iopub.status.busy":"2023-11-08T04:15:56.761563Z","iopub.execute_input":"2023-11-08T04:15:56.761957Z","iopub.status.idle":"2023-11-08T04:17:38.820903Z","shell.execute_reply.started":"2023-11-08T04:15:56.761924Z","shell.execute_reply":"2023-11-08T04:17:38.818999Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"!ls ./ -a","metadata":{"id":"ZXCUqpoaOJM5","outputId":"2e3e4913-0297-48c6-8acf-6058302080f4","execution":{"iopub.status.busy":"2023-11-08T04:17:38.823899Z","iopub.execute_input":"2023-11-08T04:17:38.825499Z","iopub.status.idle":"2023-11-08T04:17:39.241724Z","shell.execute_reply.started":"2023-11-08T04:17:38.825433Z","shell.execute_reply":"2023-11-08T04:17:39.240398Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# [Info] TF-Hub Publications\n\nAfter we saved the model in **`SavedModel`** format, we can start working on publicaiton process in [TF-Hub](https://www.tensorflow.org/hub) if we want. Basically, its a **three** step process, and [documented here](https://www.tensorflow.org/hub/publish#overview_of_the_publishing_process) in great details. Mainly,\n\n- [Export Model](https://www.tensorflow.org/hub/exporting_tf2_saved_model)\n- [Write Documentation](https://www.tensorflow.org/hub/writing_documentation)\n- [Send PR](https://www.tensorflow.org/hub/contribute_a_model).\n\n\nNote, at the time of exporting the model, you might want to exclude the optimizer, shown below. So that, you won't need to use `custom_objects` while reloading model later.\n\n```python\n\nmodel.save('saved_model', include_optimizer=False)\nreconstructed_model = keras.models.load_model(\n    'saved_model'\n)\n\n```","metadata":{"id":"oZiMt_J-OJM5"}},{"cell_type":"markdown","source":"# Conclusion\n\nIn this experiment, we've tried to inspect the visual attributes of a hybrid model with gradcam technique. In a results, it appears that the transformers blocks have more potential to interpret its decision making process with strong visual maps. The current findings are promising. We also observed that mid-level features maps of CNN tend to rectify more globally by the transformer blocks. \n\nThe visual maps created by the transformer, as shown here, are encouraging. However, it is not robust; it sometimes fails to capture relevant features. However, we could well rethink the integration strategy and develop a more effective hybrid model. Further training on more extensive datasets will generalize this hybrid model. You can try this model in [hugging face spaces](https://huggingface.co/spaces/innat/HybridModel-GradCAM). ","metadata":{"id":"dU8jjMx7OJM5"}},{"cell_type":"markdown","source":"### References\n\n- [Swin Transformer](https://arxiv.org/pdf/2103.14030.pdf)\n- [VcampSoldiers/Swin-Transformer-Tensorflow](https://github.com/VcampSoldiers/Swin-Transformer-Tensorflow)","metadata":{"id":"94OAl4M2OJM5"}},{"cell_type":"code","source":"tflite_model_path = \"/kaggle/working/model.tflite\"\ninterpreter = tf.lite.Interpreter(model_path='/kaggle/working/model.tflite')\ninterpreter.allocate_tensors()\n\ninput_details = interpreter.get_input_details()\n\ninput_detail = input_details[0]\n\ninput_data_type = input_detail['dtype']\ninput_data_shape = input_detail['shape']\n\nprint(f\"Input data type: {input_data_type}\")\nprint(f\"Input data shape: {input_data_shape}\")\n\noutput_details = interpreter.get_output_details()\n# print(output_details)\noutput_detail = input_details[0]\n\noutput_data_type = output_detail['dtype']\noutput_data_shape = output_detail['shape']\nprint(f\"Output data type : {output_data_type}\")\nprint(f\"Output data shape: {output_data_shape}\")","metadata":{"execution":{"iopub.status.busy":"2023-11-08T04:17:39.243558Z","iopub.execute_input":"2023-11-08T04:17:39.24391Z","iopub.status.idle":"2023-11-08T04:17:39.265387Z","shell.execute_reply.started":"2023-11-08T04:17:39.243872Z","shell.execute_reply":"2023-11-08T04:17:39.263812Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# import tensorflow as tf\n# import numpy as np\n# from PIL import Image\n\n# # Load the TFLite model (already done)\n# tflite_model_path = \"/kaggle/working/model.tflite\"\n# interpreter = tf.lite.Interpreter(model_path=tflite_model_path)\n# interpreter.allocate_tensors()\n\n# # Prepare the input image (replace 'your_image.jpg' with your image file)\n# input_image_path = '/kaggle/input/flowers-recognition/flowers/daisy/100080576_f52e8ee070_n.jpg'\n# input_image = Image.open(input_image_path)\n# # input_image = input_image.resize((384, 384))  # Resize to match the model's input shape\n# # input_image = np.a\n# # img_arrays = next(iter(val_ds))[0]\n# # print(img_arrays.shape)\n\n# input_image = get_img_array(input_image)\n# # input_image = np.array(input_image) / 255.0  # Normalize the pixel values to [0, 1]\n# # input_image = np.expand_dims(input_image, axis=0)  # Add batch dimension\n\n# # Get the input and output details\n# input_details = interpreter.get_input_details()\n# output_details = interpreter.get_output_details()\n\n# # Set the input tensor and run inference\n# interpreter.set_tensor(input_details[0]['index'], input_image)\n# interpreter.invoke()\n# # Get the output tensor\n# output_tensor = interpreter.get_tensor(output_details[0]['index'])\n\n# predicted_class = np.argmax(output_tensor)","metadata":{"execution":{"iopub.status.busy":"2023-11-08T04:17:39.266979Z","iopub.execute_input":"2023-11-08T04:17:39.267487Z","iopub.status.idle":"2023-11-08T04:17:39.88725Z","shell.execute_reply.started":"2023-11-08T04:17:39.267451Z","shell.execute_reply":"2023-11-08T04:17:39.886162Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"tflite_model_path = \"/kaggle/working/model.tflite\"\ninterpreter = tf.lite.Interpreter(model_path=tflite_model_path)\ninterpreter.allocate_tensors()\n\n# Define a function to preprocess the input image\ndef preprocess_input_image(image_path, input_shape=(384, 384)):\n    input_image = Image.open(image_path)\n    input_image = input_image.resize(input_shape)  # Resize to match the model's input shape\n    input_image = np.array(input_image, dtype=np.float32) / 255.0  # Normalize the pixel values to [0, 1]\n    input_image = np.expand_dims(input_image, axis=0)  # Add batch dimension\n\n    return input_image\n\n# Get the input and output details\ninput_details = interpreter.get_input_details()\noutput_details = interpreter.get_output_details()\n\n# Prepare and preprocess the input image\ninput_image = preprocess_input_image('/kaggle/input/flowers-recognition/flowers/tulip/10995953955_089572caf0.jpg')\n\n# Set the input tensor and run inference\ninterpreter.set_tensor(input_details[0]['index'], input_image)\ninterpreter.invoke()\n\n# Get the output tensor\noutput_tensor = interpreter.get_tensor(output_details[0]['index'])\n\npredicted_class = np.argmax(output_tensor)\n\n# ['daisy', 'dandelion', 'roses', 'sunflowers', 'tulips']\nclass_labels = tcls_names\n\npredicted_class_label = class_labels[predicted_class]\nprint(f\"Predicted class: {predicted_class_label}\")","metadata":{"execution":{"iopub.status.busy":"2023-11-08T04:35:03.752041Z","iopub.execute_input":"2023-11-08T04:35:03.75253Z","iopub.status.idle":"2023-11-08T04:35:04.052244Z","shell.execute_reply.started":"2023-11-08T04:35:03.752501Z","shell.execute_reply":"2023-11-08T04:35:04.051308Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"model.predict(input_image)","metadata":{"execution":{"iopub.status.busy":"2023-11-08T04:29:43.686362Z","iopub.execute_input":"2023-11-08T04:29:43.686908Z","iopub.status.idle":"2023-11-08T04:29:43.902389Z","shell.execute_reply.started":"2023-11-08T04:29:43.686872Z","shell.execute_reply":"2023-11-08T04:29:43.900894Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"reconstructed_model.predict(input_image)","metadata":{"execution":{"iopub.status.busy":"2023-11-08T04:29:54.402519Z","iopub.execute_input":"2023-11-08T04:29:54.403063Z","iopub.status.idle":"2023-11-08T04:29:54.768802Z","shell.execute_reply.started":"2023-11-08T04:29:54.403021Z","shell.execute_reply":"2023-11-08T04:29:54.76772Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Initialize the webcam or camera\ncap = cv2.VideoCapture(0)  # Use 0 for the default camera, or provide the camera index\n\nwhile True:\n    # Capture a frame from the camera\n    ret, frame = cap.read()\n\n    if not ret:\n        break\n\n    # Preprocess the input frame\n    input_image = preprocess_input_image(frame)\n\n    # Set the input tensor and run inference\n    interpreter.set_tensor(input_details[0]['index'], input_image)\n    interpreter.invoke()\n\n    # Get the output tensor\n    output_tensor = interpreter.get_tensor(output_details[0]['index'])\n\n    # Perform post-processing on the output tensor if needed\n\n    # Assuming the output is class probabilities, you can get the class with the highest probability\n    predicted_class = np.argmax(output_tensor)\n\n    class_labels = tcls_names\n\n    predicted_class_label = class_labels[predicted_class]\n\n    # Display the frame with the predicted class label\n    cv2.putText(frame, f\"Predicted class: {predicted_class_label}\", (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2)\n    cv2.imshow(\"Real-time Prediction\", frame)\n\n    # Exit the loop when the 'q' key is pressed\n    if cv2.waitKey(1) & 0xFF == ord('q'):\n        break\n\n# Release the camera and close the OpenCV windows\ncap.release()\ncv2.destroyAllWindows()","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{},"execution_count":null,"outputs":[]}]}