title a vision transformer without attention aritra roy gosthipatyhttps twitter comarig23498 ritwik rahahttps twitter comritwikraha converted to keras 3 muhammad anas razahttps anasrz com date created 20220224 last modified 20230715 description a minimal implementation of shiftvit accelerator gpu introduction vision transformershttps arxiv orgabs2010 11929 vits have sparked a wave of research at the intersection of transformers and computer vision cv vits can simultaneously model long and shortrange dependencies thanks to the multihead selfattention mechanism in the transformer block many researchers believe that the success of vits are purely due to the attention layer and they seldom think about other parts of the vit model in the academic paper when shift operation meets vision transformer an extremely simple alternative to attention mechanismhttps arxiv orgabs2201 10801 the s propose to demystify the success of vits with the introduction of a no parameter operation in place of the attention operation they swap the attention operation with a shifting operation in this example we minimally implement the paper with close alignement to the s official implementationhttps github commicrosoftspachblobmainmodelsshiftvit py setup and imports setting seed for reproducibiltiy hyperparameters these are the hyperparameters that we have chosen for the experiment please feel free to tune them data augmentation architecture optimizer training load the cifar10 dataset we use the cifar10 dataset for our experiments data augmentation the augmentation pipeline consists of rescaling resizing random cropping random horizontal flipping note the image data augmentation layers do not apply data transformations at inference time this means that when these layers are called with trainingfalse they behave differently refer to the documentationhttps keras ioapilayerspreprocessinglayersimageaugmentation for more details build the data augmentation model dataaugmentation keras sequential layers resizing config inputshape0 20 config inputshape0 20 layers randomcropconfig imagesize config imagesize layers randomfliphorizontal layers rescaling1 255 0 return dataaugmentation the shiftvit architecture the shiftvit block the mlp block class mlplayers layer def initself mlpexpandratio mlpdropoutrate kwargs super initkwargs self mlpexpandratio mlpexpandratio self mlpdropoutrate mlpdropoutrate def buildself inputshape inputchannels inputshape1 initialfilters intself mlpexpandratio inputchannels self mlp keras sequential layers dense unitsinitialfilters activationtf nn gelu layers dropoutrateself mlpdropoutrate layers denseunitsinputchannels layers dropoutrateself mlpdropoutrate def callself x x self mlpx return x the droppath layer class droppathlayers layer def initself droppathprob kwargs super initkwargs self droppathprob droppathprob def callself x trainingfalse if training keepprob 1 self droppathprob shape tf shapex0 1 lenx shape 1 randomtensor keepprob tf random uniformshape 0 1 randomtensor tf floorrandomtensor return x keepprob randomtensor return x block class shiftvitblocklayers layer def init self epsilon droppathprob mlpdropoutrate numdiv12 shiftpixel1 mlpexpandratio2 kwargs super initkwargs self shiftpixel shiftpixel self mlpexpandratio mlpexpandratio self mlpdropoutrate mlpdropoutrate self numdiv numdiv self epsilon epsilon self droppathprob droppathprob def buildself inputshape self h inputshape1 self w inputshape2 self c inputshape3 self layernorm layers layernormalizationepsilonself epsilon self droppath droppathdroppathprobself droppathprob if self droppathprob 0 0 else layers activationlinear self mlp mlp mlpexpandratioself mlpexpandratio mlpdropoutrateself mlpdropoutrate def getshiftpadself x mode split the feature maps shift the feature maps concatenate the shifted and unshifted feature maps add the residual connection the shiftvit blocks shift blokcshttps i imgur comfky5nnd png figure 5 shift blocks in the architecture sourcehttps arxiv orgabs2201 10801 each stage of the architecture has shift blocks as shown in fig 5 each of these blocks contain a variable number of stacked shiftvit block as built in the earlier section shift blocks are followed by a patchmerging layer that scales down feature inputs the patchmerging layer helps in the pyramidal structure of the model the patchmerging layer this layer merges the two adjacent tokens this layer helps in scaling the features down spatially and increasing the features up channel wise we use a conv2d layer to merge the patches the patch merging layer args epsilon float the epsilon constant apply the patch merging algorithm on the feature maps stacked shift blocks each stage will have a variable number of stacked shiftvit blocks as suggested in the paper this is a generic layer that will contain the stacked shift vit blocks with the patch merging layer as well combining the two operations shift vit block and patch merging is a design choice we picked for better code reusability note this layer will have a different depth of stacking for different stages on the model the layer containing stacked shiftvitblocks args epsilon float the epsilon constant mlpdropoutrate float the dropout rate used in the mlp block numshiftblocks int the number of shift vit blocks for this stage stochasticdepthrate float the maximum drop path rate chosen ismerge boolean a flag that determines the use of the patch merge layer after the shift vit blocks numdiv int the division of channels of the feature map defaults to 12 shiftpixel int the number of pixels to shift defaults to 1 mlpexpandratio int the ratio with which the initial dense layer of the mlp is expanded defaults to 2 calculate stochastic depth probabilities reference https keras ioexamplesvisioncctthefinalcctmodel build the shift blocks as a list of shiftvit blocks the shiftvit model build the shiftvit custom model the shiftvit model args dataaugmentation keras model a data augmentation model projecteddim int the dimension to which the patches of the image are projected patchsize int the patch size of the images numshiftblocksperstages listint a list of all the number of shit blocks per stage epsilon float the epsilon constant mlpdropoutrate float the dropout rate used in the mlp block stochasticdepthrate float the maximum drop rate probability numdiv int the number of divisions of the channesl of the feature map defaults to 12 shiftpixel int the number of pixel to shift defaults to 1 mlpexpandratio int the ratio with which the initial mlp dense layer is expanded to defaults to 2 this is the last stage do not use the patch merge here build the stages augment the images create patches and project the pathces pass through the stages get the logits calculate the loss and return it apply gradients optimize the gradients update the metrics return a dict mapping metric names to current value update the metrics return a dict mapping metric names to current value instantiate the model learning rate schedule in many experiments we want to warm up the model with a slowly increasing learning rate and then cool down the model with a slowly decaying learning rate in the warmup cosine decay the learning rate linearly increases for the warmup steps and then decays with a cosine decay some code is taken from https www kaggle comashusmatrainingrfcxtensorflowtpueffnetb2 a learningrateschedule that uses a warmup cosine decay schedule def initself lrstart lrmax warmupsteps totalsteps super init self lrstart lrstart self lrmax lrmax self warmupsteps warmupsteps self totalsteps totalsteps self pi tf constantnp pi def callself step check whether the total number of steps is larger than the warmup steps if not then throw a value error if self totalsteps self warmupsteps raise valueerror ftotal number of steps self totalsteps must be flarger or equal to warmup steps self warmupsteps cosannealedlr is a graph that increases to 1 from the initial step to the warmup step after that this graph decays to 1 at the final step mark cosannealedlr tf cos self pi tf caststep tf float32 self warmupsteps tf castself totalsteps self warmupsteps tf float32 shift the mean of the cosannealedlr graph to 1 now the grpah goes from 0 to 2 normalize the graph with 0 5 so that now it goes from 0 to 1 with the normalized graph we scale it with lrmax such that it goes from 0 to lrmax learningrate 0 5 self lrmax 1 cosannealedlr check whether warmupsteps is more than 0 if self warmupsteps 0 check whether lrmax is larger that lrstart if not throw a value error if self lrmax self lrstart raise valueerror flrstart self lrstart must be smaller or fequal to lrmax self lrmax calculate the slope with which the learning rate should increase in the warumup schedule the formula for slope is m basteps slope self lrmax self lrstart self warmupsteps with the formula for a straight line y mxc build the warmup schedule warmuprate slope tf caststep tf float32 self lrstart when the current step is lesser that warmup steps get the line graph when the current step is greater than the warmup steps get the scaled cos graph learningrate tf where step self warmupsteps warmuprate learningrate when the current step is more that the total steps return 0 else return the calculated graph return tf where step self totalsteps 0 0 learningrate namelearningrate compile and train the model get the total number of steps for training totalsteps intlenxtrain config batchsize config epochs calculate the number of steps for warmup warmupepochpercentage 0 15 warmupsteps inttotalsteps warmupepochpercentage initialize the warmupcosine schedule scheduledlrs warmupcosine lrstart1e5 lrmax1e3 warmupstepswarmupsteps totalstepstotalsteps get the optimizer optimizer keras optimizers adamw learningratescheduledlrs weightdecayconfig weightdecay compile and pretrain the model model compile optimizeroptimizer losskeras losses sparsecategoricalcrossentropyfromlogitstrue metrics keras metrics sparsecategoricalaccuracynameaccuracy keras metrics sparsetopkcategoricalaccuracy5 nametop5accuracy train the model history model fit trainds epochsconfig epochs validationdatavalds callbacks keras callbacks earlystopping monitorvalaccuracy patience5 modeauto evaluate the model with the test dataset printtesting loss acctop1 acctop5 model evaluatetestds printfloss loss 0 2f printftop 1 test accuracy acctop1100 0 2f printftop 5 test accuracy acctop5100 0 2f conclusion introduction vision transformers https arxiv org abs 2010 11929 vits have sparked a wave of research at the intersection of transformers and computer vision cv vits can simultaneously model long and short range dependencies thanks to the multi head self attention mechanism in the transformer block many researchers believe that the success of vits are purely due to the attention layer and they seldom think about other parts of the vit model in the academic paper when shift operation meets vision transformer an extremely simple alternative to attention mechanism https arxiv org abs 2201 10801 the s propose to demystify the success of vits with the introduction of a no parameter operation in place of the attention operation they swap the attention operation with a shifting operation in this example we minimally implement the paper with close alignement to the s official implementation https github com microsoft spach blob main models shiftvit py setup and imports setting seed for reproducibiltiy hyperparameters these are the hyperparameters that we have chosen for the experiment please feel free to tune them data augmentation architecture optimizer training load the cifar 10 dataset we use the cifar 10 dataset for our experiments data augmentation the augmentation pipeline consists of rescaling resizing random cropping random horizontal flipping _note_ the image data augmentation layers do not apply data transformations at inference time this means that when these layers are called with training false they behave differently refer to the documentation https keras io api layers preprocessing_layers image_augmentation for more details build the data augmentation model the shiftvit architecture in this section we build the architecture proposed in the shiftvit paper https arxiv org abs 2201 10801 shiftvit architecture https i imgur com chu40hx png figure 1 the entire architecutre of shiftvit source https arxiv org abs 2201 10801 the architecture as shown in fig 1 is inspired by swin transformer hierarchical vision transformer using shifted windows https arxiv org abs 2103 14030 here the s propose a modular architecture with 4 stages each stage works on its own spatial size creating a hierarchical architecture an input image of size hxwx3 is split into non overlapping patches of size 4x4 this is done via the patchify layer which results in individual tokens of feature size 48 4x4x3 each stage comprises two parts 1 embedding generation 2 stacked shift blocks we discuss the stages and the modules in detail in what follows _note_ compared to the official implementation https github com microsoft spach blob main models shiftvit py we restructure some key components to better fit the keras api the shiftvit block shiftvit block https i imgur com ide35vo gif figure 2 from the model to a shift block each stage in the shiftvit architecture comprises of a shift block as shown in fig 2 shift vit block https i imgur com 0q13plu png figure 3 the shift vit block source https arxiv org abs 2201 10801 the shift block as shown in fig 3 comprises of the following 1 shift operation 2 linear normalization 3 mlp layer the mlp block the mlp block is intended to be a stack of densely connected layers s get the mlp layer for each shift block args mlp_expand_ratio int the ratio with which the first feature map is expanded mlp_dropout_rate float the rate for dropout the droppath layer stochastic depth is a regularization technique that randomly drops a set of layers during inference the layers are kept as they are it is very similar to dropout but it operates on a block of layers rather than on individual nodes present inside a layer drop path also known as the stochastic depth layer refernece https keras io examples vision cct stochastic depth for regularization github com rwightman pytorch image models block the most important operation in this paper is the shift opperation in this section we describe the shift operation and compare it with its original implementation provided by the s a generic feature map is assumed to have the shape n h w c here we choose a num_div parameter that decides the division size of the channels the first 4 divisions are shifted 1 pixel in the left right up and down direction the remaining splits are kept as is after partial shifting the shifted channels are padded and the overflown pixels are chopped off this completes the partial shifting operation in the original implementation the code is approximately python out g 0 g 1 1 x g 0 g 1 1 shift left out g 1 g 2 1 x g 1 g 2 1 shift right out g 2 g 3 1 x g 2 g 3 1 shift up out g 3 g 4 1 x g 3 g 4 1 shift down out g 4 x g 4 no shift in tensorflow it would be infeasible for us to assign shifted channels to a tensor in the middle of the training process this is why we have resorted to the following procedure 1 split the channels with the num_div parameter 2 select each of the first four spilts and shift and pad them in the respective directions 3 after shifting and padding we concatenate the channel back manim rendered animation for shift operation https i imgur com preeulp gif figure 4 the tensorflow style shifting the entire procedure is explained in the fig 4 a unit shiftvit block args shift_pixel int the number of pixels to shift defaults to 1 mlp_expand_ratio int the ratio with which mlp features are expanded defaults to 2 mlp_dropout_rate float the dropout rate used in mlp num_div int the number of divisions of the feature map s channel totally 4 num_div of channels will be shifted defaults to 12 epsilon float epsilon constant drop_path_prob float the drop probability for drop path shifts the channels according to the mode chosen split the feature maps shift the feature maps concatenate the shifted and unshifted feature maps add the residual connection the shiftvit blocks shift blokcs https i imgur com fky5nnd png figure 5 shift blocks in the architecture source https arxiv org abs 2201 10801 each stage of the architecture has shift blocks as shown in fig 5 each of these blocks contain a variable number of stacked shiftvit block as built in the earlier section shift blocks are followed by a patchmerging layer that scales down feature inputs the patchmerging layer helps in the pyramidal structure of the model the patchmerging layer this layer merges the two adjacent tokens this layer helps in scaling the features down spatially and increasing the features up channel wise we use a conv2d layer to merge the patches the patch merging layer args epsilon float the epsilon constant apply the patch merging algorithm on the feature maps stacked shift blocks each stage will have a variable number of stacked shiftvit blocks as suggested in the paper this is a generic layer that will contain the stacked shift vit blocks with the patch merging layer as well combining the two operations shift vit block and patch merging is a design choice we picked for better code reusability note this layer will have a different depth of stacking for different stages on the model the layer containing stacked shiftvitblocks args epsilon float the epsilon constant mlp_dropout_rate float the dropout rate used in the mlp block num_shift_blocks int the number of shift vit blocks for this stage stochastic_depth_rate float the maximum drop path rate chosen is_merge boolean a flag that determines the use of the patch merge layer after the shift vit blocks num_div int the division of channels of the feature map defaults to 12 shift_pixel int the number of pixels to shift defaults to 1 mlp_expand_ratio int the ratio with which the initial dense layer of the mlp is expanded defaults to 2 calculate stochastic depth probabilities reference https keras io examples vision cct the final cct model build the shift blocks as a list of shiftvit blocks the shiftvit model build the shiftvit custom model the shiftvit model args data_augmentation keras model a data augmentation model projected_dim int the dimension to which the patches of the image are projected patch_size int the patch size of the images num_shift_blocks_per_stages list int a list of all the number of shit blocks per stage epsilon float the epsilon constant mlp_dropout_rate float the dropout rate used in the mlp block stochastic_depth_rate float the maximum drop rate probability num_div int the number of divisions of the channesl of the feature map defaults to 12 shift_pixel int the number of pixel to shift defaults to 1 mlp_expand_ratio int the ratio with which the initial mlp dense layer is expanded to defaults to 2 this is the last stage do not use the patch merge here build the stages augment the images create patches and project the pathces pass through the stages get the logits calculate the loss and return it apply gradients optimize the gradients update the metrics return a dict mapping metric names to current value update the metrics return a dict mapping metric names to current value instantiate the model learning rate schedule in many experiments we want to warm up the model with a slowly increasing learning rate and then cool down the model with a slowly decaying learning rate in the warmup cosine decay the learning rate linearly increases for the warmup steps and then decays with a cosine decay some code is taken from https www kaggle com ashusma training rfcx tensorflow tpu effnet b2 a learningrateschedule that uses a warmup cosine decay schedule args lr_start the initial learning rate lr_max the maximum learning rate to which lr should increase to in the warmup steps warmup_steps the number of steps for which the model warms up total_steps the total number of steps for the model training check whether the total number of steps is larger than the warmup steps if not then throw a value error cos_annealed_lr is a graph that increases to 1 from the initial step to the warmup step after that this graph decays to 1 at the final step mark shift the mean of the cos_annealed_lr graph to 1 now the grpah goes from 0 to 2 normalize the graph with 0 5 so that now it goes from 0 to 1 with the normalized graph we scale it with lr_max such that it goes from 0 to lr_max check whether warmup_steps is more than 0 check whether lr_max is larger that lr_start if not throw a value error calculate the slope with which the learning rate should increase in the warumup schedule the formula for slope is m b a steps with the formula for a straight line y mx c build the warmup schedule when the current step is lesser that warmup steps get the line graph when the current step is greater than the warmup steps get the scaled cos graph when the current step is more that the total steps return 0 else return the calculated graph compile and train the model get the total number of steps for training calculate the number of steps for warmup initialize the warmupcosine schedule get the optimizer compile and pretrain the model train the model evaluate the model with the test dataset conclusion the most impactful contribution of the paper is not the novel architecture but the idea that hierarchical vits trained with no attention can perform quite well this opens up the question of how essential attention is to the performance of vits for curious minds we would suggest reading the convnext https arxiv org abs 2201 03545 paper which attends more to the training paradigms and architectural details of vits rather than providing a novel architecture based on attention acknowledgements we would like to thank pyimagesearch https pyimagesearch com for providing us with resources that helped in the completion of this project we would like to thank jarvislabs ai https jarvislabs ai for providing with the gpu credits we would like to thank manim community https www manim community for the manim library a personal note of thanks to puja roychowdhury https twitter com pleb_talks for helping us with the learning rate schedule
import os

os.environ["KERAS_BACKEND"] = "tensorflow"

import numpy as np
import matplotlib.pyplot as plt

import tensorflow as tf
import keras
from keras import layers



SEED = 42
keras.utils.set_random_seed(SEED)




class Config(object):
    
    batch_size = 256
    buffer_size = batch_size * 2
    input_shape = (32, 32, 3)
    num_classes = 10

    
    image_size = 48

    
    patch_size = 4
    projected_dim = 96
    num_shift_blocks_per_stages = [2, 4, 8, 2]
    epsilon = 1e-5
    stochastic_depth_rate = 0.2
    mlp_dropout_rate = 0.2
    num_div = 12
    shift_pixel = 1
    mlp_expand_ratio = 2

    
    lr_start = 1e-5
    lr_max = 1e-3
    weight_decay = 1e-4

    
    epochs = 100


config = Config()



(x_train, y_train), (x_test, y_test) = keras.datasets.cifar10.load_data()
(x_train, y_train), (x_val, y_val) = (
    (x_train[:40000], y_train[:40000]),
    (x_train[40000:], y_train[40000:]),
)
print(f"Training samples: {len(x_train)}")
print(f"Validation samples: {len(x_val)}")
print(f"Testing samples: {len(x_test)}")

AUTO = tf.data.AUTOTUNE
train_ds = tf.data.Dataset.from_tensor_slices((x_train, y_train))
train_ds = (
    train_ds.shuffle(config.buffer_size).batch(config.batch_size).prefetch(AUTO)
)

val_ds = tf.data.Dataset.from_tensor_slices((x_val, y_val))
val_ds = val_ds.batch(config.batch_size).prefetch(AUTO)

test_ds = tf.data.Dataset.from_tensor_slices((x_test, y_test))
test_ds = test_ds.batch(config.batch_size).prefetch(AUTO)




def get_augmentation_model():
    
    data_augmentation = keras.Sequential(
        [
            layers.Resizing(
                config.input_shape[0] + 20, config.input_shape[0] + 20
            ),
            layers.RandomCrop(config.image_size, config.image_size),
            layers.RandomFlip("horizontal"),
            layers.Rescaling(1 / 255.0),
        ]
    )
    return data_augmentation









class MLP(layers.Layer):
    

    def __init__(self, mlp_expand_ratio, mlp_dropout_rate, **kwargs):
        super().__init__(**kwargs)
        self.mlp_expand_ratio = mlp_expand_ratio
        self.mlp_dropout_rate = mlp_dropout_rate

    def build(self, input_shape):
        input_channels = input_shape[-1]
        initial_filters = int(self.mlp_expand_ratio * input_channels)

        self.mlp = keras.Sequential(
            [
                layers.Dense(
                    units=initial_filters,
                    activation=tf.nn.gelu,
                ),
                layers.Dropout(rate=self.mlp_dropout_rate),
                layers.Dense(units=input_channels),
                layers.Dropout(rate=self.mlp_dropout_rate),
            ]
        )

    def call(self, x):
        x = self.mlp(x)
        return x





class DropPath(layers.Layer):
    

    def __init__(self, drop_path_prob, **kwargs):
        super().__init__(**kwargs)
        self.drop_path_prob = drop_path_prob

    def call(self, x, training=False):
        if training:
            keep_prob = 1 - self.drop_path_prob
            shape = (tf.shape(x)[0],) + (1,) * (len(x.shape) - 1)
            random_tensor = keep_prob + tf.random.uniform(shape, 0, 1)
            random_tensor = tf.floor(random_tensor)
            return (x / keep_prob) * random_tensor
        return x





class ShiftViTBlock(layers.Layer):
    

    def __init__(
        self,
        epsilon,
        drop_path_prob,
        mlp_dropout_rate,
        num_div=12,
        shift_pixel=1,
        mlp_expand_ratio=2,
        **kwargs,
    ):
        super().__init__(**kwargs)
        self.shift_pixel = shift_pixel
        self.mlp_expand_ratio = mlp_expand_ratio
        self.mlp_dropout_rate = mlp_dropout_rate
        self.num_div = num_div
        self.epsilon = epsilon
        self.drop_path_prob = drop_path_prob

    def build(self, input_shape):
        self.H = input_shape[1]
        self.W = input_shape[2]
        self.C = input_shape[3]
        self.layer_norm = layers.LayerNormalization(epsilon=self.epsilon)
        self.drop_path = (
            DropPath(drop_path_prob=self.drop_path_prob)
            if self.drop_path_prob > 0.0
            else layers.Activation("linear")
        )
        self.mlp = MLP(
            mlp_expand_ratio=self.mlp_expand_ratio,
            mlp_dropout_rate=self.mlp_dropout_rate,
        )

    def get_shift_pad(self, x, mode):
        
        if mode == "left":
            offset_height = 0
            offset_width = 0
            target_height = 0
            target_width = self.shift_pixel
        elif mode == "right":
            offset_height = 0
            offset_width = self.shift_pixel
            target_height = 0
            target_width = self.shift_pixel
        elif mode == "up":
            offset_height = 0
            offset_width = 0
            target_height = self.shift_pixel
            target_width = 0
        else:
            offset_height = self.shift_pixel
            offset_width = 0
            target_height = self.shift_pixel
            target_width = 0
        crop = tf.image.crop_to_bounding_box(
            x,
            offset_height=offset_height,
            offset_width=offset_width,
            target_height=self.H - target_height,
            target_width=self.W - target_width,
        )
        shift_pad = tf.image.pad_to_bounding_box(
            crop,
            offset_height=offset_height,
            offset_width=offset_width,
            target_height=self.H,
            target_width=self.W,
        )
        return shift_pad

    def call(self, x, training=False):
        
        x_splits = tf.split(
            x, num_or_size_splits=self.C // self.num_div, axis=-1
        )

        
        x_splits[0] = self.get_shift_pad(x_splits[0], mode="left")
        x_splits[1] = self.get_shift_pad(x_splits[1], mode="right")
        x_splits[2] = self.get_shift_pad(x_splits[2], mode="up")
        x_splits[3] = self.get_shift_pad(x_splits[3], mode="down")

        
        x = tf.concat(x_splits, axis=-1)

        
        shortcut = x
        x = shortcut + self.drop_path(
            self.mlp(self.layer_norm(x)), training=training
        )
        return x







class PatchMerging(layers.Layer):
    

    def __init__(self, epsilon, **kwargs):
        super().__init__(**kwargs)
        self.epsilon = epsilon

    def build(self, input_shape):
        filters = 2 * input_shape[-1]
        self.reduction = layers.Conv2D(
            filters=filters,
            kernel_size=2,
            strides=2,
            padding="same",
            use_bias=False,
        )
        self.layer_norm = layers.LayerNormalization(epsilon=self.epsilon)

    def call(self, x):
        
        x = self.layer_norm(x)
        x = self.reduction(x)
        return x







class StackedShiftBlocks(layers.Layer):
    

    def __init__(
        self,
        epsilon,
        mlp_dropout_rate,
        num_shift_blocks,
        stochastic_depth_rate,
        is_merge,
        num_div=12,
        shift_pixel=1,
        mlp_expand_ratio=2,
        **kwargs,
    ):
        super().__init__(**kwargs)
        self.epsilon = epsilon
        self.mlp_dropout_rate = mlp_dropout_rate
        self.num_shift_blocks = num_shift_blocks
        self.stochastic_depth_rate = stochastic_depth_rate
        self.is_merge = is_merge
        self.num_div = num_div
        self.shift_pixel = shift_pixel
        self.mlp_expand_ratio = mlp_expand_ratio

    def build(self, input_shapes):
        
        
        dpr = [
            x
            for x in np.linspace(
                start=0,
                stop=self.stochastic_depth_rate,
                num=self.num_shift_blocks,
            )
        ]

        
        self.shift_blocks = list()
        for num in range(self.num_shift_blocks):
            self.shift_blocks.append(
                ShiftViTBlock(
                    num_div=self.num_div,
                    epsilon=self.epsilon,
                    drop_path_prob=dpr[num],
                    mlp_dropout_rate=self.mlp_dropout_rate,
                    shift_pixel=self.shift_pixel,
                    mlp_expand_ratio=self.mlp_expand_ratio,
                )
            )
        if self.is_merge:
            self.patch_merge = PatchMerging(epsilon=self.epsilon)

    def call(self, x, training=False):
        for shift_block in self.shift_blocks:
            x = shift_block(x, training=training)
        if self.is_merge:
            x = self.patch_merge(x)
        return x





class ShiftViTModel(keras.Model):
    

    def __init__(
        self,
        data_augmentation,
        projected_dim,
        patch_size,
        num_shift_blocks_per_stages,
        epsilon,
        mlp_dropout_rate,
        stochastic_depth_rate,
        num_div=12,
        shift_pixel=1,
        mlp_expand_ratio=2,
        **kwargs,
    ):
        super().__init__(**kwargs)
        self.data_augmentation = data_augmentation
        self.patch_projection = layers.Conv2D(
            filters=projected_dim,
            kernel_size=patch_size,
            strides=patch_size,
            padding="same",
        )
        self.stages = list()
        for index, num_shift_blocks in enumerate(num_shift_blocks_per_stages):
            if index == len(num_shift_blocks_per_stages) - 1:
                
                is_merge = False
            else:
                is_merge = True
            
            self.stages.append(
                StackedShiftBlocks(
                    epsilon=epsilon,
                    mlp_dropout_rate=mlp_dropout_rate,
                    num_shift_blocks=num_shift_blocks,
                    stochastic_depth_rate=stochastic_depth_rate,
                    is_merge=is_merge,
                    num_div=num_div,
                    shift_pixel=shift_pixel,
                    mlp_expand_ratio=mlp_expand_ratio,
                )
            )
        self.global_avg_pool = layers.GlobalAveragePooling2D()

    def get_config(self):
        config = super().get_config()
        config.update(
            {
                "data_augmentation": self.data_augmentation,
                "patch_projection": self.patch_projection,
                "stages": self.stages,
                "global_avg_pool": self.global_avg_pool,
            }
        )
        return config

    def _calculate_loss(self, data, training=False):
        (images, labels) = data

        
        augmented_images = self.data_augmentation(images, training=training)

        
        projected_patches = self.patch_projection(augmented_images)

        
        x = projected_patches
        for stage in self.stages:
            x = stage(x, training=training)

        
        logits = self.global_avg_pool(x)

        
        total_loss = self.compute_loss(data, labels, logits)
        return total_loss, labels, logits

    def train_step(self, inputs):
        with tf.GradientTape() as tape:
            total_loss, labels, logits = self._calculate_loss(
                data=inputs, training=True
            )

        
        train_vars = [
            self.data_augmentation.trainable_variables,
            self.patch_projection.trainable_variables,
            self.global_avg_pool.trainable_variables,
        ]
        train_vars = train_vars + [
            stage.trainable_variables for stage in self.stages
        ]

        
        grads = tape.gradient(total_loss, train_vars)
        trainable_variable_list = []
        for grad, var in zip(grads, train_vars):
            for g, v in zip(grad, var):
                trainable_variable_list.append((g, v))
        self.optimizer.apply_gradients(trainable_variable_list)

        
        for metric in self.metrics:
            if metric.name == "loss":
                metric.update_state(total_loss)
            else:
                metric.update_state(labels, logits)

        
        return {m.name: m.result() for m in self.metrics}

    def test_step(self, data):
        loss, labels, logits = self._calculate_loss(data=data, training=False)

        
        for metric in self.metrics:
            if metric.name == "loss":
                metric.update_state(loss)
            else:
                metric.update_state(labels, logits)

        
        return {m.name: m.result() for m in self.metrics}




model = ShiftViTModel(
    data_augmentation=get_augmentation_model(),
    projected_dim=config.projected_dim,
    patch_size=config.patch_size,
    num_shift_blocks_per_stages=config.num_shift_blocks_per_stages,
    epsilon=config.epsilon,
    mlp_dropout_rate=config.mlp_dropout_rate,
    stochastic_depth_rate=config.stochastic_depth_rate,
    num_div=config.num_div,
    shift_pixel=config.shift_pixel,
    mlp_expand_ratio=config.mlp_expand_ratio,
)






class WarmUpCosine(keras.optimizers.schedules.LearningRateSchedule):
    

    def __init__(self, lr_start, lr_max, warmup_steps, total_steps):
        
        super().__init__()
        self.lr_start = lr_start
        self.lr_max = lr_max
        self.warmup_steps = warmup_steps
        self.total_steps = total_steps
        self.pi = tf.constant(np.pi)

    def __call__(self, step):
        
        
        if self.total_steps < self.warmup_steps:
            raise ValueError(
                f"Total number of steps {self.total_steps} must be"
                + f"larger or equal to warmup steps {self.warmup_steps}."
            )

        
        
        
        cos_annealed_lr = tf.cos(
            self.pi
            * (tf.cast(step, tf.float32) - self.warmup_steps)
            / tf.cast(self.total_steps - self.warmup_steps, tf.float32)
        )

        
        
        
        
        learning_rate = 0.5 * self.lr_max * (1 + cos_annealed_lr)

        
        if self.warmup_steps > 0:
            
            
            if self.lr_max < self.lr_start:
                raise ValueError(
                    f"lr_start {self.lr_start} must be smaller or"
                    + f"equal to lr_max {self.lr_max}."
                )

            
            
            slope = (self.lr_max - self.lr_start) / self.warmup_steps

            
            
            warmup_rate = slope * tf.cast(step, tf.float32) + self.lr_start

            
            
            
            learning_rate = tf.where(
                step < self.warmup_steps, warmup_rate, learning_rate
            )

        
        
        return tf.where(
            step > self.total_steps, 0.0, learning_rate, name="learning_rate"
        )





total_steps = int((len(x_train) / config.batch_size) * config.epochs)


warmup_epoch_percentage = 0.15
warmup_steps = int(total_steps * warmup_epoch_percentage)


scheduled_lrs = WarmUpCosine(
    lr_start=1e-5,
    lr_max=1e-3,
    warmup_steps=warmup_steps,
    total_steps=total_steps,
)


optimizer = keras.optimizers.AdamW(
    learning_rate=scheduled_lrs, weight_decay=config.weight_decay
)


model.compile(
    optimizer=optimizer,
    loss=keras.losses.SparseCategoricalCrossentropy(from_logits=True),
    metrics=[
        keras.metrics.SparseCategoricalAccuracy(name="accuracy"),
        keras.metrics.SparseTopKCategoricalAccuracy(5, name="top-5-accuracy"),
    ],
)


history = model.fit(
    train_ds,
    epochs=config.epochs,
    validation_data=val_ds,
    callbacks=[
        keras.callbacks.EarlyStopping(
            monitor="val_accuracy",
            patience=5,
            mode="auto",
        )
    ],
)


print("TESTING")
loss, acc_top1, acc_top5 = model.evaluate(test_ds)
print(f"Loss: {loss:0.2f}")
print(f"Top 1 test accuracy: {acc_top1*100:0.2f}%")
print(f"Top 5 test accuracy: {acc_top5*100:0.2f}%")