CINXE.COM

<!DOCTYPE html> <html lang="en"> <head> <meta charset="utf-8"> <meta name="viewport" content="width=device-width, initial-scale=1"> <meta name="description" content="Keras documentation"> <meta name="author" content="Keras Team"> <link rel="shortcut icon" href="https://keras.io/img/favicon.ico"> <link rel="canonical" href="https://keras.io/examples/vision/image_classification_using_global_context_vision_transformer/" />  <meta property="og:title" content="Keras documentation: Image Classification using Global Context Vision Transformer"> <meta property="og:image" content="https://keras.io/img/logo-k-keras-wb.png"> <meta name="twitter:title" content="Keras documentation: Image Classification using Global Context Vision Transformer"> <meta name="twitter:image" content="https://keras.io/img/k-keras-social.png"> <meta name="twitter:card" content="summary"> <title>Image Classification using Global Context Vision Transformer</title>  <link href="/css/bootstrap.min.css" rel="stylesheet">  <link href="https://fonts.googleapis.com/css2?family=Open+Sans:wght@400;600;700;800&display=swap" rel="stylesheet">  <link href="/css/docs.css" rel="stylesheet"> <link href="/css/monokai.css" rel="stylesheet">  <script>(function(w,d,s,l,i){w[l]=w[l]||[];w[l].push({'gtm.start': new Date().getTime(),event:'gtm.js'});var f=d.getElementsByTagName(s)[0], j=d.createElement(s),dl=l!='dataLayer'?'&l='+l:'';j.async=true;j.src= 'https://www.googletagmanager.com/gtm.js?id='+i+dl;f.parentNode.insertBefore(j,f); })(window,document,'script','dataLayer','GTM-5DNGF4N'); </script> <script> (function(i,s,o,g,r,a,m){i['GoogleAnalyticsObject']=r;i[r]=i[r]||function(){ (i[r].q=i[r].q||[]).push(arguments)},i[r].l=1*new Date();a=s.createElement(o), m=s.getElementsByTagName(o)[0];a.async=1;a.src=g;m.parentNode.insertBefore(a,m) })(window,document,'script','https://www.google-analytics.com/analytics.js','ga'); ga('create', 'UA-175165319-128', 'auto'); ga('send', 'pageview'); </script>  <script async defer src="https://buttons.github.io/buttons.js"></script> </head> <body>  <noscript><iframe src="https://www.googletagmanager.com/ns.html?id=GTM-5DNGF4N" height="0" width="0" style="display:none;visibility:hidden"></iframe></noscript>  <div class='k-page'> <div class="k-nav" id="nav-menu"> <a href='/'><img src='/img/logo-small.png' class='logo-small' /></a> <div class="nav flex-column nav-pills" role="tablist" aria-orientation="vertical"> <a class="nav-link" href="/about/" role="tab" aria-selected="">About Keras</a> <a class="nav-link" href="/getting_started/" role="tab" aria-selected="">Getting started</a> <a class="nav-link" href="/guides/" role="tab" aria-selected="">Developer guides</a> <a class="nav-link active" href="/examples/" role="tab" aria-selected="">Code examples</a> <a class="nav-sublink active" href="/examples/vision/">Computer Vision</a> <a class="nav-sublink2" href="/examples/vision/image_classification_from_scratch/">Image classification from scratch</a> <a class="nav-sublink2" href="/examples/vision/mnist_convnet/">Simple MNIST convnet</a> <a class="nav-sublink2" href="/examples/vision/image_classification_efficientnet_fine_tuning/">Image classification via fine-tuning with EfficientNet</a> <a class="nav-sublink2" href="/examples/vision/image_classification_with_vision_transformer/">Image classification with Vision Transformer</a> <a class="nav-sublink2" href="/examples/vision/attention_mil_classification/">Classification using Attention-based Deep Multiple Instance Learning</a> <a class="nav-sublink2" href="/examples/vision/mlp_image_classification/">Image classification with modern MLP models</a> <a class="nav-sublink2" href="/examples/vision/mobilevit/">A mobile-friendly Transformer-based model for image classification</a> <a class="nav-sublink2" href="/examples/vision/xray_classification_with_tpus/">Pneumonia Classification on TPU</a> <a class="nav-sublink2" href="/examples/vision/cct/">Compact Convolutional Transformers</a> <a class="nav-sublink2" href="/examples/vision/convmixer/">Image classification with ConvMixer</a> <a class="nav-sublink2" href="/examples/vision/eanet/">Image classification with EANet (External Attention Transformer)</a> <a class="nav-sublink2" href="/examples/vision/involution/">Involutional neural networks</a> <a class="nav-sublink2" href="/examples/vision/perceiver_image_classification/">Image classification with Perceiver</a> <a class="nav-sublink2" href="/examples/vision/reptile/">Few-Shot learning with Reptile</a> <a class="nav-sublink2" href="/examples/vision/semisupervised_simclr/">Semi-supervised image classification using contrastive pretraining with SimCLR</a> <a class="nav-sublink2" href="/examples/vision/swin_transformers/">Image classification with Swin Transformers</a> <a class="nav-sublink2" href="/examples/vision/vit_small_ds/">Train a Vision Transformer on small datasets</a> <a class="nav-sublink2" href="/examples/vision/shiftvit/">A Vision Transformer without Attention</a> <a class="nav-sublink2 active" href="/examples/vision/image_classification_using_global_context_vision_transformer/">Image Classification using Global Context Vision Transformer</a> <a class="nav-sublink2" href="/examples/vision/oxford_pets_image_segmentation/">Image segmentation with a U-Net-like architecture</a> <a class="nav-sublink2" href="/examples/vision/deeplabv3_plus/">Multiclass semantic segmentation using DeepLabV3+</a> <a class="nav-sublink2" href="/examples/vision/basnet_segmentation/">Highly accurate boundaries segmentation using BASNet</a> <a class="nav-sublink2" href="/examples/vision/fully_convolutional_network/">Image Segmentation using Composable Fully-Convolutional Networks</a> <a class="nav-sublink2" href="/examples/vision/retinanet/">Object Detection with RetinaNet</a> <a class="nav-sublink2" href="/examples/vision/keypoint_detection/">Keypoint Detection with Transfer Learning</a> <a class="nav-sublink2" href="/examples/vision/object_detection_using_vision_transformer/">Object detection with Vision Transformers</a> <a class="nav-sublink2" href="/examples/vision/3D_image_classification/">3D image classification from CT scans</a> <a class="nav-sublink2" href="/examples/vision/depth_estimation/">Monocular depth estimation</a> <a class="nav-sublink2" href="/examples/vision/nerf/">3D volumetric rendering with NeRF</a> <a class="nav-sublink2" href="/examples/vision/pointnet_segmentation/">Point cloud segmentation with PointNet</a> <a class="nav-sublink2" href="/examples/vision/pointnet/">Point cloud classification</a> <a class="nav-sublink2" href="/examples/vision/captcha_ocr/">OCR model for reading Captchas</a> <a class="nav-sublink2" href="/examples/vision/handwriting_recognition/">Handwriting recognition</a> <a class="nav-sublink2" href="/examples/vision/autoencoder/">Convolutional autoencoder for image denoising</a> <a class="nav-sublink2" href="/examples/vision/mirnet/">Low-light image enhancement using MIRNet</a> <a class="nav-sublink2" href="/examples/vision/super_resolution_sub_pixel/">Image Super-Resolution using an Efficient Sub-Pixel CNN</a> <a class="nav-sublink2" href="/examples/vision/edsr/">Enhanced Deep Residual Networks for single-image super-resolution</a> <a class="nav-sublink2" href="/examples/vision/zero_dce/">Zero-DCE for low-light image enhancement</a> <a class="nav-sublink2" href="/examples/vision/cutmix/">CutMix data augmentation for image classification</a> <a class="nav-sublink2" href="/examples/vision/mixup/">MixUp augmentation for image classification</a> <a class="nav-sublink2" href="/examples/vision/randaugment/">RandAugment for Image Classification for Improved Robustness</a> <a class="nav-sublink2" href="/examples/vision/image_captioning/">Image captioning</a> <a class="nav-sublink2" href="/examples/vision/nl_image_search/">Natural language image search with a Dual Encoder</a> <a class="nav-sublink2" href="/examples/vision/visualizing_what_convnets_learn/">Visualizing what convnets learn</a> <a class="nav-sublink2" href="/examples/vision/integrated_gradients/">Model interpretability with Integrated Gradients</a> <a class="nav-sublink2" href="/examples/vision/probing_vits/">Investigating Vision Transformer representations</a> <a class="nav-sublink2" href="/examples/vision/grad_cam/">Grad-CAM class activation visualization</a> <a class="nav-sublink2" href="/examples/vision/near_dup_search/">Near-duplicate image search</a> <a class="nav-sublink2" href="/examples/vision/semantic_image_clustering/">Semantic Image Clustering</a> <a class="nav-sublink2" href="/examples/vision/siamese_contrastive/">Image similarity estimation using a Siamese Network with a contrastive loss</a> <a class="nav-sublink2" href="/examples/vision/siamese_network/">Image similarity estimation using a Siamese Network with a triplet loss</a> <a class="nav-sublink2" href="/examples/vision/metric_learning/">Metric learning for image similarity search</a> <a class="nav-sublink2" href="/examples/vision/metric_learning_tf_similarity/">Metric learning for image similarity search using TensorFlow Similarity</a> <a class="nav-sublink2" href="/examples/vision/nnclr/">Self-supervised contrastive learning with NNCLR</a> <a class="nav-sublink2" href="/examples/vision/video_classification/">Video Classification with a CNN-RNN Architecture</a> <a class="nav-sublink2" href="/examples/vision/conv_lstm/">Next-Frame Video Prediction with Convolutional LSTMs</a> <a class="nav-sublink2" href="/examples/vision/video_transformers/">Video Classification with Transformers</a> <a class="nav-sublink2" href="/examples/vision/vivit/">Video Vision Transformer</a> <a class="nav-sublink2" href="/examples/vision/bit/">Image Classification using BigTransfer (BiT)</a> <a class="nav-sublink2" href="/examples/vision/gradient_centralization/">Gradient Centralization for Better Training Performance</a> <a class="nav-sublink2" href="/examples/vision/token_learner/">Learning to tokenize in Vision Transformers</a> <a class="nav-sublink2" href="/examples/vision/knowledge_distillation/">Knowledge Distillation</a> <a class="nav-sublink2" href="/examples/vision/fixres/">FixRes: Fixing train-test resolution discrepancy</a> <a class="nav-sublink2" href="/examples/vision/cait/">Class Attention Image Transformers with LayerScale</a> <a class="nav-sublink2" href="/examples/vision/patch_convnet/">Augmenting convnets with aggregated attention</a> <a class="nav-sublink2" href="/examples/vision/learnable_resizer/">Learning to Resize</a> <a class="nav-sublink2" href="/examples/vision/adamatch/">Semi-supervision and domain adaptation with AdaMatch</a> <a class="nav-sublink2" href="/examples/vision/barlow_twins/">Barlow Twins for Contrastive SSL</a> <a class="nav-sublink2" href="/examples/vision/consistency_training/">Consistency training with supervision</a> <a class="nav-sublink2" href="/examples/vision/deit/">Distilling Vision Transformers</a> <a class="nav-sublink2" href="/examples/vision/focal_modulation_network/">Focal Modulation: A replacement for Self-Attention</a> <a class="nav-sublink2" href="/examples/vision/forwardforward/">Using the Forward-Forward Algorithm for Image Classification</a> <a class="nav-sublink2" href="/examples/vision/masked_image_modeling/">Masked image modeling with Autoencoders</a> <a class="nav-sublink2" href="/examples/vision/sam/">Segment Anything Model with 🤗Transformers</a> <a class="nav-sublink2" href="/examples/vision/segformer/">Semantic segmentation with SegFormer and Hugging Face Transformers</a> <a class="nav-sublink2" href="/examples/vision/simsiam/">Self-supervised contrastive learning with SimSiam</a> <a class="nav-sublink2" href="/examples/vision/supervised-contrastive-learning/">Supervised Contrastive Learning</a> <a class="nav-sublink2" href="/examples/vision/temporal_latent_bottleneck/">When Recurrence meets Transformers</a> <a class="nav-sublink2" href="/examples/vision/yolov8/">Efficient Object Detection with YOLOV8 and KerasCV</a> <a class="nav-sublink" href="/examples/nlp/">Natural Language Processing</a> <a class="nav-sublink" href="/examples/structured_data/">Structured Data</a> <a class="nav-sublink" href="/examples/timeseries/">Timeseries</a> <a class="nav-sublink" href="/examples/generative/">Generative Deep Learning</a> <a class="nav-sublink" href="/examples/audio/">Audio Data</a> <a class="nav-sublink" href="/examples/rl/">Reinforcement Learning</a> <a class="nav-sublink" href="/examples/graph/">Graph Data</a> <a class="nav-sublink" href="/examples/keras_recipes/">Quick Keras Recipes</a> <a class="nav-link" href="/api/" role="tab" aria-selected="">Keras 3 API documentation</a> <a class="nav-link" href="/2.18/api/" role="tab" aria-selected="">Keras 2 API documentation</a> <a class="nav-link" href="/keras_tuner/" role="tab" aria-selected="">KerasTuner: Hyperparam Tuning</a> <a class="nav-link" href="/keras_hub/" role="tab" aria-selected="">KerasHub: Pretrained Models</a> </div> </div> <div class='k-main'> <div class='k-main-top'> <script> function displayDropdownMenu() { e = document.getElementById("nav-menu"); if (e.style.display == "block") { e.style.display = "none"; } else { e.style.display = "block"; document.getElementById("dropdown-nav").style.display = "block"; } } function resetMobileUI() { if (window.innerWidth <= 840) { document.getElementById("nav-menu").style.display = "none"; document.getElementById("dropdown-nav").style.display = "block"; } else { document.getElementById("nav-menu").style.display = "block"; document.getElementById("dropdown-nav").style.display = "none"; } var navmenu = document.getElementById("nav-menu"); var menuheight = navmenu.clientHeight; var kmain = document.getElementById("k-main-id"); kmain.style.minHeight = (menuheight + 100) + 'px'; } window.onresize = resetMobileUI; window.addEventListener("load", (event) => { resetMobileUI() }); </script> <div id='dropdown-nav' onclick="displayDropdownMenu();"> <svg viewBox="-20 -20 120 120" width="60" height="60"> <rect width="100" height="20"></rect> <rect y="30" width="100" height="20"></rect> <rect y="60" width="100" height="20"></rect> </svg> </div> <form class="bd-search d-flex align-items-center k-search-form" id="search-form"> <input type="search" class="k-search-input" id="search-input" placeholder="Search Keras documentation..." aria-label="Search Keras documentation..." autocomplete="off"> <button class="k-search-btn"> <svg width="13" height="13" viewBox="0 0 13 13"><title>search</title><path d="m4.8495 7.8226c0.82666 0 1.5262-0.29146 2.0985-0.87438 0.57232-0.58292 0.86378-1.2877 0.87438-2.1144 0.010599-0.82666-0.28086-1.5262-0.87438-2.0985-0.59352-0.57232-1.293-0.86378-2.0985-0.87438-0.8055-0.010599-1.5103 0.28086-2.1144 0.87438-0.60414 0.59352-0.8956 1.293-0.87438 2.0985 0.021197 0.8055 0.31266 1.5103 0.87438 2.1144 0.56172 0.60414 1.2665 0.8956 2.1144 0.87438zm4.4695 0.2115 3.681 3.6819-1.259 1.284-3.6817-3.7 0.0019784-0.69479-0.090043-0.098846c-0.87973 0.76087-1.92 1.1413-3.1207 1.1413-1.3553 0-2.5025-0.46363-3.4417-1.3909s-1.4088-2.0686-1.4088-3.4239c0-1.3553 0.4696-2.4966 1.4088-3.4239 0.9392-0.92727 2.0864-1.3969 3.4417-1.4088 1.3553-0.011889 2.4906 0.45771 3.406 1.4088 0.9154 0.95107 1.379 2.0924 1.3909 3.4239 0 1.2126-0.38043 2.2588-1.1413 3.1385l0.098834 0.090049z"></path></svg> </button> </form> <script> var form = document.getElementById('search-form'); form.onsubmit = function(e) { e.preventDefault(); var query = document.getElementById('search-input').value; window.location.href = '/search.html?query=' + query; return False } </script> </div> <div class='k-main-inner' id='k-main-id'> <div class='k-location-slug'> ► <a href='/examples/'>Code examples</a> / <a href='/examples/vision/'>Computer Vision</a> / Image Classification using Global Context Vision Transformer </div> <div class='k-content'> <h1 id="image-classification-using-global-context-vision-transformer">Image Classification using Global Context Vision Transformer</h1> Author: Md Awsafur Rahman Date created: 2023/10/30 Last modified: 2023/10/30 Description: Implementation and fine-tuning of Global Context Vision Transformer for image classification. <div class='example_version_banner keras_3'>ⓘ This example uses Keras 3</div> <img class="k-inline-icon" src="https://colab.research.google.com/img/colab_favicon.ico"/> <a href="https://colab.research.google.com/github/keras-team/keras-io/blob/master/examples/vision/ipynb/image_classification_using_global_context_vision_transformer.ipynb">View in Colab</a> •<img class="k-inline-icon" src="https://github.com/favicon.ico"/> <a href="https://github.com/keras-team/keras-io/blob/master/examples/vision/image_classification_using_global_context_vision_transformer.py">GitHub source</a> <h1 id="setup">Setup</h1> <div class="codehilite"><pre><code>!pip install --upgrade keras_cv tensorflow !pip install --upgrade keras </code></pre></div> <div class="codehilite"><pre><code>import keras from keras_cv.layers import DropPath from keras import ops from keras import layers import tensorflow as tf # only for dataloader import tensorflow_datasets as tfds # for flower dataset from skimage.data import chelsea import matplotlib.pyplot as plt import numpy as np </code></pre></div> <hr /> <h2 id="introduction">Introduction</h2> In this notebook, we will utilize multi-backend Keras 3.0 to implement the <a href="https://arxiv.org/abs/2206.09959">GCViT: Global Context Vision Transformer</a> paper, presented at ICML 2023 by A Hatamizadeh et al. The, we will fine-tune the model on the Flower dataset for image classification task, leveraging the official ImageNet pre-trained weights. A highlight of this notebook is its compatibility with multiple backends: TensorFlow, PyTorch, and JAX, showcasing the true potential of multi-backend Keras. <hr /> <h2 id="motivation">Motivation</h2> <blockquote> Note: In this section we'll learn about the backstory of GCViT and try to understand why it is proposed. </blockquote> <ul> <li>During recent years, Transformers have achieved dominance in Natural Language Processing (NLP) tasks and with the self-attention mechanism which allows for capturing both long and short-range information.</li> <li>Following this trend, Vision Transformer (ViT) proposed to utilize image patches as tokens in a gigantic architecture similar to encoder of the original Transformer.</li> <li>Despite the historic dominance of Convolutional Neural Network (CNN) in computer vision, ViT-based models have shown SOTA or competitive performance in various computer vision tasks.</li> </ul> <img src="https://raw.githubusercontent.com/awsaf49/gcvit-tf/main/image/vit_gif.gif" width=600> <ul> <li>However, quadratic [<code>O(n^2)</code>] computational complexity of self-attention and lack of multi-scale information makes it difficult for ViT to be considered as general-purpose architecture for Compute Vision tasks like segmentation and object detection where it requires dense prediction at the pixel level.</li> <li>Swin Transformer has attempted to address the issues of ViT by proposing multi-resolution/hierarchical architectures in which the self-attention is computed in local windows and cross-window connections such as window shifting are used for modeling the interactions across different regions. But the limited receptive field of local windows can not capture long-range information, and cross-window-connection schemes such as window-shifting only cover a small neighborhood in the vicinity of each window. Also, it lacks inductive-bias that encourages certain translation invariance is still preferable for general-purpose visual modeling, particularly for the dense prediction tasks of object detection and semantic segmentation.</li> </ul> <img src="https://raw.githubusercontent.com/awsaf49/gcvit-tf/main/image/swin_vs_vit.JPG" width=400> <img src="https://raw.githubusercontent.com/awsaf49/gcvit-tf/main/image/shifted_window.JPG" width=400> <img src="https://raw.githubusercontent.com/awsaf49/gcvit-tf/main/image/swin_arch.JPG" width=800> <ul> <li>To address above limitations, Global Context (GC) ViT network is proposed.</li> </ul> <hr /> <h2 id="architecture">Architecture</h2> Let's have a quick overview of our key components, 1. <code>Stem/PatchEmbed:</code> A stem/patchify layer processes images at the network’s beginning. For this network, it creates patches/tokens and converts them into embeddings. 2. <code>Level:</code> It is the repetitive building block that extracts features using different blocks. 3. <code>Global Token Gen./FeatureExtraction:</code> It generates global tokens/patches with Depthwise-CNN, SqueezeAndExcitation (Squeeze-Excitation), CNN and MaxPooling. So basically it's a Feature Extractor. 4. <code>Block:</code> It is the repetitive module that applies attention to the features and projects them to a certain dimension. 1. <code>Local-MSA:</code> Local Multi head Self Attention. 2. <code>Global-MSA:</code> Global Multi head Self Attention. 3. <code>MLP:</code> Linear layer that projects a vector to another dimension. 5. <code>Downsample/ReduceSize:</code> It is very similar to Global Token Gen. module except it uses CNN instead of MaxPooling to downsample with additional Layer Normalization modules. 6. <code>Head:</code> It is the module responsible for the classification task. 1. <code>Pooling:</code> It converts <code>N x 2D</code> features to <code>N x 1D</code> features. 2. <code>Classifier:</code> It processes <code>N x 1D</code> features to make a decision about class. I've annotated the architecture figure to make it easier to digest, <img src="https://raw.githubusercontent.com/awsaf49/gcvit-tf/main/image/arch_annot.png"> <h3 id="unit-blocks">Unit Blocks</h3> <blockquote> Note: This blocks are used to build other modules throughout the paper. Most of the blocks are either borrowed from other work or modified version old work. </blockquote> <ol> <li> <code>SqueezeAndExcitation</code>: Squeeze-Excitation (SE) aka Bottleneck module acts sd kind of channel attention. It consits of AvgPooling, Dense/FullyConnected (FC)/Linear , GELU and Sigmoid module. <img src="https://raw.githubusercontent.com/awsaf49/gcvit-tf/main/image/se_annot.png" width=400> </li> <li> <code>Fused-MBConv:</code> This is similar to the one used in EfficientNetV2. It uses Depthwise-Conv, GELU, SqueezeAndExcitation, Conv, to extract feature with a resiudal connection. Note that, no new module is declared for this one, we simply applied corresponding modules directly. <img src="https://raw.githubusercontent.com/awsaf49/gcvit-tf/main/image/fmb_annot.png" width=350> </li> <li> <code>ReduceSize</code>: It is a CNN based downsample module which abvobe mentioned <code>Fused-MBConv</code> module to extract feature, Strided Conv to simultaneously reduce spatial dimension and increse channelwise dimention of the features and finally LayerNormalization module to normalize features. In the paper/figure this module is referred as downsample module. I think it is mention worthy that SwniTransformer used <code>PatchMerging</code> module instead of <code>ReduceSize</code> to reduce the spatial dimention and increase channelwise dimension which uses fully-connected/dense/linear module. According to the GCViT paper, one of the purposes of using <code>ReduceSize</code> is to add inductive bias through CNN module. <img src="https://raw.githubusercontent.com/awsaf49/gcvit-tf/main/image/down_annot.png" width=300> </li> <li> <code>MLP:</code> This is our very own Multi Layer Perceptron module. This a feed-forward/fully-connected/linear module which simply projects input to an arbitary dimension. </li> </ol> <div class="codehilite"><pre><code>class SqueezeAndExcitation(layers.Layer): """Squeeze and excitation block. Args: output_dim: output features dimension, if `None` use same dim as input. expansion: expansion ratio. """ def __init__(self, output_dim=None, expansion=0.25, **kwargs): super().__init__(**kwargs) self.expansion = expansion self.output_dim = output_dim def build(self, input_shape): inp = input_shape[-1] self.output_dim = self.output_dim or inp self.avg_pool = layers.GlobalAvgPool2D(keepdims=True, name="avg_pool") self.fc = [ layers.Dense(int(inp * self.expansion), use_bias=False, name="fc_0"), layers.Activation("gelu", name="fc_1"), layers.Dense(self.output_dim, use_bias=False, name="fc_2"), layers.Activation("sigmoid", name="fc_3"), ] super().build(input_shape) def call(self, inputs, **kwargs): x = self.avg_pool(inputs) for layer in self.fc: x = layer(x) return x * inputs class ReduceSize(layers.Layer): """Down-sampling block. Args: keepdims: if False spatial dim is reduced and channel dim is increased """ def __init__(self, keepdims=False, **kwargs): super().__init__(**kwargs) self.keepdims = keepdims def build(self, input_shape): embed_dim = input_shape[-1] dim_out = embed_dim if self.keepdims else 2 * embed_dim self.pad1 = layers.ZeroPadding2D(1, name="pad1") self.pad2 = layers.ZeroPadding2D(1, name="pad2") self.conv = [ layers.DepthwiseConv2D( kernel_size=3, strides=1, padding="valid", use_bias=False, name="conv_0" ), layers.Activation("gelu", name="conv_1"), SqueezeAndExcitation(name="conv_2"), layers.Conv2D( embed_dim, kernel_size=1, strides=1, padding="valid", use_bias=False, name="conv_3", ), ] self.reduction = layers.Conv2D( dim_out, kernel_size=3, strides=2, padding="valid", use_bias=False, name="reduction", ) self.norm1 = layers.LayerNormalization( -1, 1e-05, name="norm1" ) # eps like PyTorch self.norm2 = layers.LayerNormalization(-1, 1e-05, name="norm2") def call(self, inputs, **kwargs): x = self.norm1(inputs) xr = self.pad1(x) for layer in self.conv: xr = layer(xr) x = x + xr x = self.pad2(x) x = self.reduction(x) x = self.norm2(x) return x class MLP(layers.Layer): """Multi-Layer Perceptron (MLP) block. Args: hidden_features: hidden features dimension. out_features: output features dimension. activation: activation function. dropout: dropout rate. """ def __init__( self, hidden_features=None, out_features=None, activation="gelu", dropout=0.0, **kwargs, ): super().__init__(**kwargs) self.hidden_features = hidden_features self.out_features = out_features self.activation = activation self.dropout = dropout def build(self, input_shape): self.in_features = input_shape[-1] self.hidden_features = self.hidden_features or self.in_features self.out_features = self.out_features or self.in_features self.fc1 = layers.Dense(self.hidden_features, name="fc1") self.act = layers.Activation(self.activation, name="act") self.fc2 = layers.Dense(self.out_features, name="fc2") self.drop1 = layers.Dropout(self.dropout, name="drop1") self.drop2 = layers.Dropout(self.dropout, name="drop2") def call(self, inputs, **kwargs): x = self.fc1(inputs) x = self.act(x) x = self.drop1(x) x = self.fc2(x) x = self.drop2(x) return x </code></pre></div> <h3 id="stem">Stem</h3> <blockquote> Notes: In the code, this module is referred to as PatchEmbed but on paper, it is referred to as Stem. </blockquote> In the model, we have first used <code>patch_embed</code> module. Let's try to understand this module. As we can see from the <code>call</code> method, 1. This module first pads input 2. Then uses convolutions to extract patches with embeddings. 3. Finally, uses <code>ReduceSize</code> module to first extract features with convolution but neither reduces spatial dimension nor increases spatial dimension. 4. One important point to notice, unlike ViT or SwinTransformer, GCViT creates overlapping patches. We can notice that from the code, <code>Conv2D(self.embed_dim, kernel_size=3, strides=2, name='proj')</code>. If we wanted non-overlapping patches then we would've used the same <code>kernel_size</code> and <code>stride</code>. 5. This module reduces the spatial dimension of input by <code>4x</code>. <blockquote> Summary: image → padding → convolution → (feature_extract + downsample) </blockquote> <div class="codehilite"><pre><code>class PatchEmbed(layers.Layer): """Patch embedding block. Args: embed_dim: feature size dimension. """ def __init__(self, embed_dim, **kwargs): super().__init__(**kwargs) self.embed_dim = embed_dim def build(self, input_shape): self.pad = layers.ZeroPadding2D(1, name="pad") self.proj = layers.Conv2D(self.embed_dim, 3, 2, name="proj") self.conv_down = ReduceSize(keepdims=True, name="conv_down") def call(self, inputs, **kwargs): x = self.pad(inputs) x = self.proj(x) x = self.conv_down(x) return x </code></pre></div> <h3 id="global-token-gen">Global Token Gen.</h3> <blockquote> Notes: It is one of the two CNN modules that is used to imppose inductive bias. </blockquote> As we can see from above cell, in the <code>level</code> we have first used <code>to_q_global/Global Token Gen./FeatureExtraction</code>. Let's try to understand how it works, <ul> <li>This module is series of <code>FeatureExtract</code> module, according to paper we need to repeat this module <code>K</code> times, where <code>K = log2(H/h)</code>, <code>H = feature_map_height</code>, <code>W = feature_map_width</code>.</li> <li><code>FeatureExtraction:</code> This layer is very similar to <code>ReduceSize</code> module except it uses MaxPooling module to reduce the dimension, it doesn't increse feature dimension (channelsie) and it doesn't uses LayerNormalizaton. This module is used to in <code>Generate Token Gen.</code> module repeatedly to generte global tokens for global-context-attention.</li> <li>One important point to notice from the figure is that, global tokens is shared across the whole image which means we use only one global window for all local tokens in a image. This makes the computation very efficient.</li> <li>For input feature map with shape <code>(B, H, W, C)</code>, we'll get output shape <code>(B, h, w, C)</code>. If we copy these global tokens for total <code>M</code> local windows in an image where, <code>M = (H x W)/(h x w) = num_window</code>, then output shape: <code>(B * M, h, w, C)</code>."</li> </ul> <blockquote> Summary: This module is used to <code>resize</code> the image to fit window. </blockquote> <img src="https://raw.githubusercontent.com/awsaf49/gcvit-tf/main/image/global_token_annot.png" width=800> <div class="codehilite"><pre><code>class FeatureExtraction(layers.Layer): """Feature extraction block. Args: keepdims: bool argument for maintaining the resolution. """ def __init__(self, keepdims=False, **kwargs): super().__init__(**kwargs) self.keepdims = keepdims def build(self, input_shape): embed_dim = input_shape[-1] self.pad1 = layers.ZeroPadding2D(1, name="pad1") self.pad2 = layers.ZeroPadding2D(1, name="pad2") self.conv = [ layers.DepthwiseConv2D(3, 1, use_bias=False, name="conv_0"), layers.Activation("gelu", name="conv_1"), SqueezeAndExcitation(name="conv_2"), layers.Conv2D(embed_dim, 1, 1, use_bias=False, name="conv_3"), ] if not self.keepdims: self.pool = layers.MaxPool2D(3, 2, name="pool") super().build(input_shape) def call(self, inputs, **kwargs): x = inputs xr = self.pad1(x) for layer in self.conv: xr = layer(xr) x = x + xr if not self.keepdims: x = self.pool(self.pad2(x)) return x class GlobalQueryGenerator(layers.Layer): """Global query generator. Args: keepdims: to keep the dimension of FeatureExtraction layer. For instance, repeating log(56/7) = 3 blocks, with input window dimension 56 and output window dimension 7 at down-sampling ratio 2. Please check Fig.5 of GC ViT paper for details. """ def __init__(self, keepdims=False, **kwargs): super().__init__(**kwargs) self.keepdims = keepdims def build(self, input_shape): self.to_q_global = [ FeatureExtraction(keepdims, name=f"to_q_global_{i}") for i, keepdims in enumerate(self.keepdims) ] super().build(input_shape) def call(self, inputs, **kwargs): x = inputs for layer in self.to_q_global: x = layer(x) return x </code></pre></div> <h3 id="attention">Attention</h3> <blockquote> Notes: This is the core contribution of the paper. </blockquote> As we can see from the <code>call</code> method, 1. <code>WindowAttention</code> module applies both local and global window attention depending on <code>global_query</code> parameter. <ol> <li>First it converts input features into <code>query, key, value</code> for local attention and <code>key, value</code> for global attention. For global attention, it takes global query from <code>Global Token Gen.</code>. One thing to notice from the code is that we divide the features or embed_dim among all the heads of Transformer to reduce the computation. <code>qkv = tf.reshape(qkv, [B_, N, self.qkv_size, self.num_heads, C // self.num_heads])</code></li> <li>Before sending query, key and value for attention, global token goes through an important process. Same global tokens or one global window gets copied for all the local windows to increase efficiency. <code>q_global = tf.repeat(q_global, repeats=B_//B, axis=0)</code>, here <code>B_//B</code> means <code>num_windows</code> in a image.</li> <li>Then simply applies <code>local-window-self-attention</code> or <code>global-window-attention</code> depending on <code>global_query</code> parameter. One thing to notice from the code is that we are adding relative-positional-embedding with the attention mask instead of the patch embedding. <code>attn = attn + relative_position_bias[tf.newaxis,]</code> <img src="https://raw.githubusercontent.com/awsaf49/gcvit-tf/main/image/lvg_msa.PNG" width=800></li> <li>Now, let's think for a bit and try to understand what is happening here. Let's focus on the figure below. We can see from the left, that in the local-attention the query is local and it's limited to the local window (red square border) hence we don't have access to long-range information. But on the right that due to global query we're now not limited to local-windows (blue square border) and we have access to long-range information. <img src="https://raw.githubusercontent.com/awsaf49/gcvit-tf/main/image/lvg_arch.PNG" width=800></li> <li>In ViT we compare (attention) image-tokens with image-tokens, in SwinTransformer we compare window-tokens with window-tokens but in GCViT we compare image-tokens with window-tokens. But now you may ask, how can compare(attention) image-tokens with window-tokens even after image-tokens have larger dimensions than window-tokens? (from above figure image-tokens have shape <code>(1, 8, 8, 3)</code> and window-tokens have shape <code>(1, 4, 4, 3)</code>). Yes, you are right we can't directly compare them hence we resize image-tokens to fit window-tokens with <code>Global Token Gen./FeatureExtraction</code> CNN module. The following table should give you a clear comparison,</li> </ol> <table> <thead> <tr> <th>Model</th> <th>Query Tokens</th> <th>Key-Value Tokens</th> <th>Attention Type</th> <th>Attention Coverage</th> </tr> </thead> <tbody> <tr> <td>ViT</td> <td>image</td> <td>image</td> <td>self-attention</td> <td>global</td> </tr> <tr> <td>SwinTransformer</td> <td>window</td> <td>window</td> <td>self-attention</td> <td>local</td> </tr> <tr> <td>GCViT</td> <td>resized-image</td> <td>window</td> <td>image-window attention</td> <td>global</td> </tr> </tbody> </table> <div class="codehilite"><pre><code>class WindowAttention(layers.Layer): """Local window attention. This implementation was proposed by [Liu et al., 2021](https://arxiv.org/abs/2103.14030) in SwinTransformer. Args: window_size: window size. num_heads: number of attention head. global_query: if the input contains global_query qkv_bias: bool argument for query, key, value learnable bias. qk_scale: bool argument to scaling query, key. attention_dropout: attention dropout rate. projection_dropout: output dropout rate. """ def __init__( self, window_size, num_heads, global_query, qkv_bias=True, qk_scale=None, attention_dropout=0.0, projection_dropout=0.0, **kwargs, ): super().__init__(**kwargs) window_size = (window_size, window_size) self.window_size = window_size self.num_heads = num_heads self.global_query = global_query self.qkv_bias = qkv_bias self.qk_scale = qk_scale self.attention_dropout = attention_dropout self.projection_dropout = projection_dropout def build(self, input_shape): embed_dim = input_shape[0][-1] head_dim = embed_dim // self.num_heads self.scale = self.qk_scale or head_dim**-0.5 self.qkv_size = 3 - int(self.global_query) self.qkv = layers.Dense( embed_dim * self.qkv_size, use_bias=self.qkv_bias, name="qkv" ) self.relative_position_bias_table = self.add_weight( name="relative_position_bias_table", shape=[ (2 * self.window_size[0] - 1) * (2 * self.window_size[1] - 1), self.num_heads, ], initializer=keras.initializers.TruncatedNormal(stddev=0.02), trainable=True, dtype=self.dtype, ) self.attn_drop = layers.Dropout(self.attention_dropout, name="attn_drop") self.proj = layers.Dense(embed_dim, name="proj") self.proj_drop = layers.Dropout(self.projection_dropout, name="proj_drop") self.softmax = layers.Activation("softmax", name="softmax") super().build(input_shape) def get_relative_position_index(self): coords_h = ops.arange(self.window_size[0]) coords_w = ops.arange(self.window_size[1]) coords = ops.stack(ops.meshgrid(coords_h, coords_w, indexing="ij"), axis=0) coords_flatten = ops.reshape(coords, [2, -1]) relative_coords = coords_flatten[:, :, None] - coords_flatten[:, None, :] relative_coords = ops.transpose(relative_coords, axes=[1, 2, 0]) relative_coords_xx = relative_coords[:, :, 0] + self.window_size[0] - 1 relative_coords_yy = relative_coords[:, :, 1] + self.window_size[1] - 1 relative_coords_xx = relative_coords_xx * (2 * self.window_size[1] - 1) relative_position_index = relative_coords_xx + relative_coords_yy return relative_position_index def call(self, inputs, **kwargs): if self.global_query: inputs, q_global = inputs B = ops.shape(q_global)[0] # B, N, C else: inputs = inputs[0] B_, N, C = ops.shape(inputs) # B*num_window, num_tokens, channels qkv = self.qkv(inputs) qkv = ops.reshape( qkv, [B_, N, self.qkv_size, self.num_heads, C // self.num_heads] ) qkv = ops.transpose(qkv, [2, 0, 3, 1, 4]) if self.global_query: k, v = ops.split( qkv, indices_or_sections=2, axis=0 ) # for unknown shame num=None will throw error q_global = ops.repeat( q_global, repeats=B_ // B, axis=0 ) # num_windows = B_//B => q_global same for all windows in a img q = ops.reshape(q_global, [B_, N, self.num_heads, C // self.num_heads]) q = ops.transpose(q, axes=[0, 2, 1, 3]) else: q, k, v = ops.split(qkv, indices_or_sections=3, axis=0) q = ops.squeeze(q, axis=0) k = ops.squeeze(k, axis=0) v = ops.squeeze(v, axis=0) q = q * self.scale attn = q @ ops.transpose(k, axes=[0, 1, 3, 2]) relative_position_bias = ops.take( self.relative_position_bias_table, ops.reshape(self.get_relative_position_index(), [-1]), ) relative_position_bias = ops.reshape( relative_position_bias, [ self.window_size[0] * self.window_size[1], self.window_size[0] * self.window_size[1], -1, ], ) relative_position_bias = ops.transpose(relative_position_bias, axes=[2, 0, 1]) attn = attn + relative_position_bias[None,] attn = self.softmax(attn) attn = self.attn_drop(attn) x = ops.transpose((attn @ v), axes=[0, 2, 1, 3]) x = ops.reshape(x, [B_, N, C]) x = self.proj_drop(self.proj(x)) return x </code></pre></div> <h3 id="block">Block</h3> <blockquote> Notes: This module doesn't have any Convolutional module. </blockquote> In the <code>level</code> second module that we have used is <code>block</code>. Let's try to understand how it works. As we can see from the <code>call</code> method, 1. <code>Block</code> module takes either only feature_maps for local attention or additional global query for global attention. 2. Before sending feature maps for attention, this module converts batch feature maps to batch windows as we'll be applying Window Attention. 3. Then we send batch batch windows for attention. 4. After attention has been applied we revert batch windows to batch feature maps. 5. Before sending the attention to applied features for output, this module applies Stochastic Depth regularization in the residual connection. Also, before applying Stochastic Depth it rescales the input with trainable parameters. Note that, this Stochastic Depth block hasn't been shown in the figure of the paper. <img src="https://raw.githubusercontent.com/awsaf49/gcvit-tf/main/image/block2.JPG" width=400> <h3 id="window">Window</h3> In the <code>block</code> module, we have created windows before and after applying attention. Let's try to understand how we're creating windows, * Following module converts feature maps <code>(B, H, W, C)</code> to stacked windows <code>(B x H/h x W/w, h, w, C)</code> → <code>(num_windows_batch, window_size, window_size, channel)</code> * This module uses <code>reshape</code> & <code>transpose</code> to create these windows out of image instead of iterating over them. <div class="codehilite"><pre><code>class Block(layers.Layer): """GCViT block. Args: window_size: window size. num_heads: number of attention head. global_query: apply global window attention mlp_ratio: MLP ratio. qkv_bias: bool argument for query, key, value learnable bias. qk_scale: bool argument to scaling query, key. drop: dropout rate. attention_dropout: attention dropout rate. path_drop: drop path rate. activation: activation function. layer_scale: layer scaling coefficient. """ def __init__( self, window_size, num_heads, global_query, mlp_ratio=4.0, qkv_bias=True, qk_scale=None, dropout=0.0, attention_dropout=0.0, path_drop=0.0, activation="gelu", layer_scale=None, **kwargs, ): super().__init__(**kwargs) self.window_size = window_size self.num_heads = num_heads self.global_query = global_query self.mlp_ratio = mlp_ratio self.qkv_bias = qkv_bias self.qk_scale = qk_scale self.dropout = dropout self.attention_dropout = attention_dropout self.path_drop = path_drop self.activation = activation self.layer_scale = layer_scale def build(self, input_shape): B, H, W, C = input_shape[0] self.norm1 = layers.LayerNormalization(-1, 1e-05, name="norm1") self.attn = WindowAttention( window_size=self.window_size, num_heads=self.num_heads, global_query=self.global_query, qkv_bias=self.qkv_bias, qk_scale=self.qk_scale, attention_dropout=self.attention_dropout, projection_dropout=self.dropout, name="attn", ) self.drop_path1 = DropPath(self.path_drop) self.drop_path2 = DropPath(self.path_drop) self.norm2 = layers.LayerNormalization(-1, 1e-05, name="norm2") self.mlp = MLP( hidden_features=int(C * self.mlp_ratio), dropout=self.dropout, activation=self.activation, name="mlp", ) if self.layer_scale is not None: self.gamma1 = self.add_weight( name="gamma1", shape=[C], initializer=keras.initializers.Constant(self.layer_scale), trainable=True, dtype=self.dtype, ) self.gamma2 = self.add_weight( name="gamma2", shape=[C], initializer=keras.initializers.Constant(self.layer_scale), trainable=True, dtype=self.dtype, ) else: self.gamma1 = 1.0 self.gamma2 = 1.0 self.num_windows = int(H // self.window_size) * int(W // self.window_size) super().build(input_shape) def call(self, inputs, **kwargs): if self.global_query: inputs, q_global = inputs else: inputs = inputs[0] B, H, W, C = ops.shape(inputs) x = self.norm1(inputs) # create windows and concat them in batch axis x = self.window_partition(x, self.window_size) # (B_, win_h, win_w, C) # flatten patch x = ops.reshape(x, [-1, self.window_size * self.window_size, C]) # attention if self.global_query: x = self.attn([x, q_global]) else: x = self.attn([x]) # reverse window partition x = self.window_reverse(x, self.window_size, H, W, C) # FFN x = inputs + self.drop_path1(x * self.gamma1) x = x + self.drop_path2(self.gamma2 * self.mlp(self.norm2(x))) return x def window_partition(self, x, window_size): """ Args: x: (B, H, W, C) window_size: window size Returns: local window features (num_windows*B, window_size, window_size, C) """ B, H, W, C = ops.shape(x) x = ops.reshape( x, [ -1, H // window_size, window_size, W // window_size, window_size, C, ], ) x = ops.transpose(x, axes=[0, 1, 3, 2, 4, 5]) windows = ops.reshape(x, [-1, window_size, window_size, C]) return windows def window_reverse(self, windows, window_size, H, W, C): """ Args: windows: local window features (num_windows*B, window_size, window_size, C) window_size: Window size H: Height of image W: Width of image C: Channel of image Returns: x: (B, H, W, C) """ x = ops.reshape( windows, [ -1, H // window_size, W // window_size, window_size, window_size, C, ], ) x = ops.transpose(x, axes=[0, 1, 3, 2, 4, 5]) x = ops.reshape(x, [-1, H, W, C]) return x </code></pre></div> <h3 id="level">Level</h3> <blockquote> Note: This module has both Transformer and CNN modules. </blockquote> In the model, the second module that we have used is <code>level</code>. Let's try to understand this module. As we can see from the <code>call</code> method, 1. First it creates global_token with a series of <code>FeatureExtraction</code> modules. As we'll see later that <code>FeatureExtraction</code> is nothing but a simple CNN based module. 2. Then it uses series of<code>Block</code> modules to apply local or global window attention depending on depth level. 3. Finally, it uses <code>ReduceSize</code> to reduce the dimension of contextualized features. <blockquote> Summary: feature_map → global_token → local/global window attention → dowsample </blockquote> <img src="https://raw.githubusercontent.com/awsaf49/gcvit-tf/main/image/level.png" width=400> <div class="codehilite"><pre><code>class Level(layers.Layer): """GCViT level. Args: depth: number of layers in each stage. num_heads: number of heads in each stage. window_size: window size in each stage. keepdims: dims to keep in FeatureExtraction. downsample: bool argument for down-sampling. mlp_ratio: MLP ratio. qkv_bias: bool argument for query, key, value learnable bias. qk_scale: bool argument to scaling query, key. drop: dropout rate. attention_dropout: attention dropout rate. path_drop: drop path rate. layer_scale: layer scaling coefficient. """ def __init__( self, depth, num_heads, window_size, keepdims, downsample=True, mlp_ratio=4.0, qkv_bias=True, qk_scale=None, dropout=0.0, attention_dropout=0.0, path_drop=0.0, layer_scale=None, **kwargs, ): super().__init__(**kwargs) self.depth = depth self.num_heads = num_heads self.window_size = window_size self.keepdims = keepdims self.downsample = downsample self.mlp_ratio = mlp_ratio self.qkv_bias = qkv_bias self.qk_scale = qk_scale self.dropout = dropout self.attention_dropout = attention_dropout self.path_drop = path_drop self.layer_scale = layer_scale def build(self, input_shape): path_drop = ( [self.path_drop] * self.depth if not isinstance(self.path_drop, list) else self.path_drop ) self.blocks = [ Block( window_size=self.window_size, num_heads=self.num_heads, global_query=bool(i % 2), mlp_ratio=self.mlp_ratio, qkv_bias=self.qkv_bias, qk_scale=self.qk_scale, dropout=self.dropout, attention_dropout=self.attention_dropout, path_drop=path_drop[i], layer_scale=self.layer_scale, name=f"blocks_{i}", ) for i in range(self.depth) ] self.down = ReduceSize(keepdims=False, name="downsample") self.q_global_gen = GlobalQueryGenerator(self.keepdims, name="q_global_gen") super().build(input_shape) def call(self, inputs, **kwargs): x = inputs q_global = self.q_global_gen(x) # shape: (B, win_size, win_size, C) for i, blk in enumerate(self.blocks): if i % 2: x = blk([x, q_global]) # shape: (B, H, W, C) else: x = blk([x]) # shape: (B, H, W, C) if self.downsample: x = self.down(x) # shape: (B, H//2, W//2, 2*C) return x </code></pre></div> <h3 id="model">Model</h3> Let's directly jump to the model. As we can see from the <code>call</code> method, 1. It creates patch embeddings from an image. This layer doesn't flattens these embeddings which means output of this module will be <code>(batch, height/window_size, width/window_size, embed_dim)</code> instead of <code>(batch, height x width/window_size^2, embed_dim)</code>. 2. Then it applies <code>Dropout</code> module which randomly sets input units to 0. 3. It passes these embeddings to series of <code>Level</code> modules which we are calling <code>level</code> where, 1. Global token is generated 1. Both local & global attention is applied 1. Finally downsample is applied. 4. So, output after <code>n</code> number of levels, shape: <code>(batch, width/window_size x 2^{n-1}, width/window_size x 2^{n-1}, embed_dim x 2^{n-1})</code>. In the last layer, paper doesn't use downsample and increase channels. 5. Output of above layer is normalized using <code>LayerNormalization</code> module. 6. In the head, 2D features are converted to 1D features with <code>Pooling</code> module. Output shape after this module is <code>(batch, embed_dim x 2^{n-1})</code> 7. Finally, pooled features are sent to <code>Dense/Linear</code> module for classification. <blockquote> Sumamry: image → (patchs + embedding) → dropout → (attention + feature extraction) → normalizaion → pooling → classify </blockquote> <div class="codehilite"><pre><code>class GCViT(keras.Model): """GCViT model. Args: window_size: window size in each stage. embed_dim: feature size dimension. depths: number of layers in each stage. num_heads: number of heads in each stage. drop_rate: dropout rate. mlp_ratio: MLP ratio. qkv_bias: bool argument for query, key, value learnable bias. qk_scale: bool argument to scaling query, key. attention_dropout: attention dropout rate. path_drop: drop path rate. layer_scale: layer scaling coefficient. num_classes: number of classes. head_activation: activation function for head. """ def __init__( self, window_size, embed_dim, depths, num_heads, drop_rate=0.0, mlp_ratio=3.0, qkv_bias=True, qk_scale=None, attention_dropout=0.0, path_drop=0.1, layer_scale=None, num_classes=1000, head_activation="softmax", **kwargs, ): super().__init__(**kwargs) self.window_size = window_size self.embed_dim = embed_dim self.depths = depths self.num_heads = num_heads self.drop_rate = drop_rate self.mlp_ratio = mlp_ratio self.qkv_bias = qkv_bias self.qk_scale = qk_scale self.attention_dropout = attention_dropout self.path_drop = path_drop self.layer_scale = layer_scale self.num_classes = num_classes self.head_activation = head_activation self.patch_embed = PatchEmbed(embed_dim=embed_dim, name="patch_embed") self.pos_drop = layers.Dropout(drop_rate, name="pos_drop") path_drops = np.linspace(0.0, path_drop, sum(depths)) keepdims = [(0, 0, 0), (0, 0), (1,), (1,)] self.levels = [] for i in range(len(depths)): path_drop = path_drops[sum(depths[:i]) : sum(depths[: i + 1])].tolist() level = Level( depth=depths[i], num_heads=num_heads[i], window_size=window_size[i], keepdims=keepdims[i], downsample=(i < len(depths) - 1), mlp_ratio=mlp_ratio, qkv_bias=qkv_bias, qk_scale=qk_scale, dropout=drop_rate, attention_dropout=attention_dropout, path_drop=path_drop, layer_scale=layer_scale, name=f"levels_{i}", ) self.levels.append(level) self.norm = layers.LayerNormalization(axis=-1, epsilon=1e-05, name="norm") self.pool = layers.GlobalAvgPool2D(name="pool") self.head = layers.Dense(num_classes, name="head", activation=head_activation) def build(self, input_shape): super().build(input_shape) self.built = True def call(self, inputs, **kwargs): x = self.patch_embed(inputs) # shape: (B, H, W, C) x = self.pos_drop(x) for level in self.levels: x = level(x) # shape: (B, H_, W_, C_) x = self.norm(x) x = self.pool(x) # shape: (B, C__) x = self.head(x) return x def build_graph(self, input_shape=(224, 224, 3)): """ ref: https://www.kaggle.com/code/ipythonx/tf-hybrid-efficientnet-swin-transformer-gradcam """ x = keras.Input(shape=input_shape) return keras.Model(inputs=[x], outputs=self.call(x), name=self.name) def summary(self, input_shape=(224, 224, 3)): return self.build_graph(input_shape).summary() </code></pre></div> <hr /> <h2 id="build-model">Build Model</h2> <ul> <li>Let's build a complete model with all the modules that we've explained above. We'll build GCViT-XXTiny model with the configuration mentioned in the paper.</li> <li>Also we'll load the ported official pre-trained weights and try for some predictions.</li> </ul> <div class="codehilite"><pre><code># Model Configs config = { "window_size": (7, 7, 14, 7), "embed_dim": 64, "depths": (2, 2, 6, 2), "num_heads": (2, 4, 8, 16), "mlp_ratio": 3.0, "path_drop": 0.2, } ckpt_link = ( "https://github.com/awsaf49/gcvit-tf/releases/download/v1.1.6/gcvitxxtiny.keras" ) # Build Model model = GCViT(**config) inp = ops.array(np.random.uniform(size=(1, 224, 224, 3))) out = model(inp) # Load Weights ckpt_path = keras.utils.get_file(ckpt_link.split("/")[-1], ckpt_link) model.load_weights(ckpt_path) # Summary model.summary((224, 224, 3)) </code></pre></div> <div class="k-default-codeblock"> <div class="codehilite"><pre><code>Downloading data from https://github.com/awsaf49/gcvit-tf/releases/download/v1.1.6/gcvitxxtiny.keras 48767519/48767519 ━━━━━━━━━━━━━━━━━━━━ 0s 0us/step </code></pre></div> </div> <pre style="white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace">Model: "gc_vi_t" </pre> <pre style="white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace">┏━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━┓ ┃ Layer (type) ┃ Output Shape ┃ Param # ┃ ┡━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━┩ │ input_layer (InputLayer) │ (None, 224, 224, 3) │ 0 │ ├────────────────────────────────────┼───────────────────────────────┼─────────────┤ │ patch_embed (PatchEmbed) │ (None, 56, 56, 64) │ 45,632 │ ├────────────────────────────────────┼───────────────────────────────┼─────────────┤ │ pos_drop (Dropout) │ (None, 56, 56, 64) │ 0 │ ├────────────────────────────────────┼───────────────────────────────┼─────────────┤ │ levels_0 (Level) │ (None, 28, 28, 128) │ 180,964 │ ├────────────────────────────────────┼───────────────────────────────┼─────────────┤ │ levels_1 (Level) │ (None, 14, 14, 256) │ 688,456 │ ├────────────────────────────────────┼───────────────────────────────┼─────────────┤ │ levels_2 (Level) │ (None, 7, 7, 512) │ 5,170,608 │ ├────────────────────────────────────┼───────────────────────────────┼─────────────┤ │ levels_3 (Level) │ (None, 7, 7, 512) │ 5,395,744 │ ├────────────────────────────────────┼───────────────────────────────┼─────────────┤ │ norm (LayerNormalization) │ (None, 7, 7, 512) │ 1,024 │ ├────────────────────────────────────┼───────────────────────────────┼─────────────┤ │ pool (GlobalAveragePooling2D) │ (None, 512) │ 0 │ ├────────────────────────────────────┼───────────────────────────────┼─────────────┤ │ head (Dense) │ (None, 1000) │ 513,000 │ └────────────────────────────────────┴───────────────────────────────┴─────────────┘ </pre> <pre style="white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace"> Total params: 11,995,428 (45.76 MB) </pre> <pre style="white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace"> Trainable params: 11,995,428 (45.76 MB) </pre> <pre style="white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace"> Non-trainable params: 0 (0.00 B) </pre> <hr /> <h2 id="sanity-check-for-pretrained-weights">Sanity check for Pre-Trained Weights</h2> <div class="codehilite"><pre><code>img = keras.applications.imagenet_utils.preprocess_input( chelsea(), mode="torch" ) # Chelsea the cat img = ops.image.resize(img, (224, 224))[None,] # resize & create batch pred = model(img) pred_dec = keras.applications.imagenet_utils.decode_predictions(pred)[0] print("\n# Image:") plt.figure(figsize=(6, 6)) plt.imshow(chelsea()) plt.show() print() print("# Prediction (Top 5):") for i in range(5): print("{:<12} : {:0.2f}".format(pred_dec[i][1], pred_dec[i][2])) </code></pre></div> <div class="k-default-codeblock"> <div class="codehilite"><pre><code>Downloading data from https://storage.googleapis.com/download.tensorflow.org/data/imagenet_class_index.json 35363/35363 ━━━━━━━━━━━━━━━━━━━━ 0s 0us/step </code></pre></div> </div> <div class="k-default-codeblock"> <div class="codehilite"><pre><code># Image: </code></pre></div> </div> <img alt="png" src="/img/examples/vision/image_classification_using_global_context_vision_transformer/image_classification_using_global_context_vision_transformer_24_1.png" /> <div class="k-default-codeblock"> <div class="codehilite"><pre><code># Prediction (Top 5): Egyptian_cat : 0.72 tiger_cat : 0.04 tabby : 0.03 crossword_puzzle : 0.01 panpipe : 0.00 </code></pre></div> </div> <h1 id="gcvit">Fine-tune GCViT Model</h1> In the following cells, we will fine-tune GCViT model on Flower Dataset which consists <code>104</code> classes. <h3 id="configs">Configs</h3> <div class="codehilite"><pre><code># Model IMAGE_SIZE = (224, 224) # Hyper Params BATCH_SIZE = 32 EPOCHS = 5 # Dataset CLASSES = [ "dandelion", "daisy", "tulips", "sunflowers", "roses", ] # don't change the order # Other constants MEAN = 255 * np.array([0.485, 0.456, 0.406], dtype="float32") # imagenet mean STD = 255 * np.array([0.229, 0.224, 0.225], dtype="float32") # imagenet std AUTO = tf.data.AUTOTUNE </code></pre></div> <hr /> <h2 id="data-loader">Data Loader</h2> <div class="codehilite"><pre><code>def make_dataset(dataset: tf.data.Dataset, train: bool, image_size: int = IMAGE_SIZE): def preprocess(image, label): # for training, do augmentation if train: if tf.random.uniform(shape=[]) > 0.5: image = tf.image.flip_left_right(image) image = tf.image.resize(image, size=image_size, method="bicubic") image = (image - MEAN) / STD # normalization return image, label if train: dataset = dataset.shuffle(BATCH_SIZE * 10) return dataset.map(preprocess, AUTO).batch(BATCH_SIZE).prefetch(AUTO) </code></pre></div> <h3 id="flower-dataset">Flower Dataset</h3> <div class="codehilite"><pre><code>train_dataset, val_dataset = tfds.load( "tf_flowers", split=["train[:90%]", "train[90%:]"], as_supervised=True, try_gcs=False, # gcs_path is necessary for tpu, ) train_dataset = make_dataset(train_dataset, True) val_dataset = make_dataset(val_dataset, False) </code></pre></div> <div class="k-default-codeblock"> <div class="codehilite"><pre><code>Downloading and preparing dataset 218.21 MiB (download: 218.21 MiB, generated: 221.83 MiB, total: 440.05 MiB) to /root/tensorflow_datasets/tf_flowers/3.0.1... Dl Completed...: 0%| | 0/5 [00:00<?, ? file/s] Dataset tf_flowers downloaded and prepared to /root/tensorflow_datasets/tf_flowers/3.0.1. Subsequent calls will reuse this data. </code></pre></div> </div> <h3 id="rebuild-model-for-flower-dataset">Re-Build Model for Flower Dataset</h3> <div class="codehilite"><pre><code># Re-Build Model model = GCViT(**config, num_classes=104) inp = ops.array(np.random.uniform(size=(1, 224, 224, 3))) out = model(inp) # Load Weights ckpt_path = keras.utils.get_file(ckpt_link.split("/")[-1], ckpt_link) model.load_weights(ckpt_path, skip_mismatch=True) model.compile( loss="sparse_categorical_crossentropy", optimizer="adam", metrics=["accuracy"] ) </code></pre></div> <div class="k-default-codeblock"> <div class="codehilite"><pre><code>/usr/local/lib/python3.10/dist-packages/keras/src/saving/saving_lib.py:269: UserWarning: A total of 1 objects could not be loaded. Example error message for object <Dense name=head, built=True>: </code></pre></div> </div> <div class="k-default-codeblock"> <div class="codehilite"><pre><code>Layer 'head' expected 2 variables, but received 0 variables during loading. Expected: ['kernel', 'bias'] </code></pre></div> </div> <div class="k-default-codeblock"> <div class="codehilite"><pre><code>List of objects that could not be loaded: [<Dense name=head, built=True>] warnings.warn(msg) </code></pre></div> </div> <h3 id="training">Training</h3> <div class="codehilite"><pre><code>history = model.fit( train_dataset, validation_data=val_dataset, epochs=EPOCHS, verbose=1 ) </code></pre></div> <div class="k-default-codeblock"> <div class="codehilite"><pre><code>Epoch 1/5 104/104 ━━━━━━━━━━━━━━━━━━━━ 153s 581ms/step - accuracy: 0.5140 - loss: 1.4615 - val_accuracy: 0.8828 - val_loss: 0.3485 Epoch 2/5 104/104 ━━━━━━━━━━━━━━━━━━━━ 7s 69ms/step - accuracy: 0.8775 - loss: 0.3437 - val_accuracy: 0.8828 - val_loss: 0.3508 Epoch 3/5 104/104 ━━━━━━━━━━━━━━━━━━━━ 7s 68ms/step - accuracy: 0.8937 - loss: 0.2918 - val_accuracy: 0.9019 - val_loss: 0.2953 Epoch 4/5 104/104 ━━━━━━━━━━━━━━━━━━━━ 7s 68ms/step - accuracy: 0.9232 - loss: 0.2397 - val_accuracy: 0.9183 - val_loss: 0.2212 Epoch 5/5 104/104 ━━━━━━━━━━━━━━━━━━━━ 7s 68ms/step - accuracy: 0.9456 - loss: 0.1645 - val_accuracy: 0.9210 - val_loss: 0.2897 </code></pre></div> </div> <hr /> <h2 id="reference">Reference</h2> <ul> <li><a href="https://github.com/awsaf49/gcvit-tf">gcvit-tf - A Python library for GCViT with TF2.0</a></li> <li><a href="https://github.com/NVlabs/GCVit">gcvit - Official codebase for GCViT</a></li> </ul> </div> <div class='k-outline'> <div class='k-outline-depth-1'> <a href='#image-classification-using-global-context-vision-transformer'>Image Classification using Global Context Vision Transformer</a> </div> <div class='k-outline-depth-1'> <a href='#setup'>Setup</a> </div> <div class='k-outline-depth-2'> ◆ <a href='#introduction'>Introduction</a> </div> <div class='k-outline-depth-2'> ◆ <a href='#motivation'>Motivation</a> </div> <div class='k-outline-depth-2'> ◆ <a href='#architecture'>Architecture</a> </div> <div class='k-outline-depth-3'> <a href='#unit-blocks'>Unit Blocks</a> </div> <div class='k-outline-depth-3'> <a href='#stem'>Stem</a> </div> <div class='k-outline-depth-3'> <a href='#global-token-gen'>Global Token Gen.</a> </div> <div class='k-outline-depth-3'> <a href='#attention'>Attention</a> </div> <div class='k-outline-depth-3'> <a href='#block'>Block</a> </div> <div class='k-outline-depth-3'> <a href='#window'>Window</a> </div> <div class='k-outline-depth-3'> <a href='#level'>Level</a> </div> <div class='k-outline-depth-3'> <a href='#model'>Model</a> </div> <div class='k-outline-depth-2'> ◆ <a href='#build-model'>Build Model</a> </div> <div class='k-outline-depth-2'> ◆ <a href='#sanity-check-for-pretrained-weights'>Sanity check for Pre-Trained Weights</a> </div> <div class='k-outline-depth-1'> <a href='#finetune-gcvit-model'>Fine-tune **GCViT** Model</a> </div> <div class='k-outline-depth-3'> <a href='#configs'>Configs</a> </div> <div class='k-outline-depth-2'> ◆ <a href='#data-loader'>Data Loader</a> </div> <div class='k-outline-depth-3'> <a href='#flower-dataset'>Flower Dataset</a> </div> <div class='k-outline-depth-3'> <a href='#rebuild-model-for-flower-dataset'>Re-Build Model for Flower Dataset</a> </div> <div class='k-outline-depth-3'> <a href='#training'>Training</a> </div> <div class='k-outline-depth-2'> ◆ <a href='#reference'>Reference</a> </div> </div> </div> </div> </div> </body> <footer style="float: left; width: 100%; padding: 1em; border-top: solid 1px #bbb;"> <a href="https://policies.google.com/terms">Terms</a> | <a href="https://policies.google.com/privacy">Privacy</a> </footer> </html>