title distilling vision transformers sayak paulhttps twitter comrisingsayak date created 20220405 last modified 20230916 description distillation of vision transformers through attention accelerator gpu introduction in the original vision transformers vit paper dosovitskiy et al https arxiv orgabs2010 11929 the s concluded that to perform on par with convolutional neural networks cnns vits need to be pretrained on larger datasets the larger the better this is mainly due to the lack of inductive biases in the vit architecture unlike cnns they don t have layers that exploit locality in a followup paper steiner et al https arxiv orgabs2106 10270 the s show that it is possible to substantially improve the performance of vits with stronger regularization and longer training many groups have proposed different ways to deal with the problem of dataintensiveness of vit training one such way was shown in the dataefficient image transformers deit paper touvron et al https arxiv orgabs2012 12877 the s introduced a distillation technique that is specific to transformerbased vision models deit is among the first works to show that it s possible to train vits well without using larger datasets in this example we implement the distillation recipe proposed in deit this requires us to slightly tweak the original vit architecture and write a custom training loop to implement the distillation recipe to comfortably navigate through this example you ll be expected to know how a vit and knowledge distillation work the following are good resources in case you needed a refresher vit on keras iohttps keras ioexamplesvisionimageclassificationwithvisiontransformer knowledge distillation on keras iohttps keras ioexamplesvisionknowledgedistillation imports constants model training data you probably noticed that dropoutrate has been set 0 0 dropout has been used in the implementation to keep it complete for smaller models like the one used in this example you don t need it but for bigger models using dropout helps load the tfflowers dataset and prepare preprocessing utilities the s use an array of different augmentation techniques including mixup zhang et al https arxiv orgabs1710 09412 randaugment cubuk et al https arxiv orgabs1909 13719 and so on however to keep the example simple to work through we ll discard them resize to a bigger spatial resolution and take the random crops implementing the deit variants of vit since deit is an extension of vit it d make sense to first implement vit and then extend it to support deit s components first we ll implement a layer for stochastic depth huang et al https arxiv orgabs1603 09382 which is used in deit for regularization referred from github com rwightmanpytorchimagemodels now we ll implement the mlp and transformer blocks ffn for a transformer block iterate over the hidden units and add dense dropout for idx units in enumeratehiddenunits x layers dense units activationtf nn gelu if idx 0 else none x x layers dropoutdropoutratex return x def transformerdropprob float name str keras model layer normalization 1 multi head self attention layer 1 skip connection 1 layer normalization 2 mlp layer 1 skip connection 2 we ll now implement a vitclassifier class building on top of the components we just developed here we ll be following the original pooling strategy used in the vit paper use a class token and use the feature representations corresponding to it for classification vision transformer base class def initself kwargs super initkwargs patchify linear projection reshaping self projection keras sequential layers conv2d filtersprojectiondim kernelsizepatchsize patchsize stridespatchsize patchsize paddingvalid nameconvprojection layers reshape targetshapenumpatches projectiondim nameflattenprojection nameprojection positional embedding initshape 1 numpatches 1 projectiondim self positionalembedding tf variable tf zerosinitshape nameposembedding transformer blocks dpr x for x in tf linspace0 0 droppathrate numlayers self transformerblocks transformerdropprobdpri nameftransformerblocki for i in rangenumlayers cls token initialvalue tf zeros1 1 projectiondim self clstoken tf variable initialvalueinitialvalue trainabletrue namecls other layers self dropout layers dropoutdropoutrate self layernorm layers layernormalizationepsilonlayernormeps self head layers dense numclasses nameclassificationhead def callself inputs n tf shapeinputs0 create patches and project the patches projectedpatches self projectioninputs append class token if needed clstoken tf tileself clstoken n 1 1 clstoken tf castclstoken projectedpatches dtype projectedpatches tf concatclstoken projectedpatches axis1 add positional embeddings to the projected patches encodedpatches self positionalembedding projectedpatches b numberpatches projectiondim encodedpatches self dropoutencodedpatches iterate over the number of layers and stack up blocks of transformer for transformermodule in self transformerblocks add a transformer block encodedpatches transformermoduleencodedpatches final layer normalization representation self layernormencodedpatches pool representation encodedpatches representation 0 classification head output self headencodedpatches return output class vitdistilledvitclassifier def initself regulartrainingfalse kwargs super initkwargs self numtokens 2 self regulartraining regulartraining cls and distillation tokens positional embedding initvalue tf zeros1 1 projectiondim self disttoken tf variableinitvalue namedisttoken self positionalembedding tf variable tf zeros 1 numpatches self numtokens projectiondim nameposembedding head layers self head layers dense numclasses nameclassificationhead self headdist layers dense numclasses namedistillationhead def callself inputs trainingfalse n tf shapeinputs0 create patches and project the patches projectedpatches self projectioninputs append the tokens clstoken tf tileself clstoken n 1 1 disttoken tf tileself disttoken n 1 1 clstoken tf castclstoken projectedpatches dtype disttoken tf castdisttoken projectedpatches dtype projectedpatches tf concat clstoken disttoken projectedpatches axis1 add positional embeddings to the projected patches encodedpatches self positionalembedding projectedpatches b numberpatches projectiondim encodedpatches self dropoutencodedpatches iterate over the number of layers and stack up blocks of transformer for transformermodule in self transformerblocks add a transformer block encodedpatches transformermoduleencodedpatches final layer normalization representation self layernormencodedpatches classification heads x xdist self headrepresentation 0 self headdistrepresentation 1 if not training or self regulartraining during standard train finetune inference average the classifier predictions return x xdist 2 elif training only return separate classification predictions when training in distilled mode return x xdist deittinydistilled vitdistilled dummyinputs tf ones2 224 224 3 outputs deittinydistilleddummyinputs trainingfalse printfoutputshape outputs shape implementing the trainer class deitkeras model reference https keras ioexamplesvisionknowledgedistillation def initself student teacher kwargs super initkwargs self student student self teacher teacher self studentlosstracker keras metrics meannamestudentloss self distillationlosstracker keras metrics mean namedistillationloss self accuracy keras metrics categoricalaccuracynameaccuracy property def metricsself return self accuracy self studentlosstracker self distillationlosstracker def compile self optimizer studentlossfn distillationlossfn runeagerlyfalse jitcompilefalse super compile optimizeroptimizer runeagerlyruneagerly jitcompilejitcompile self studentlossfn studentlossfn self distillationlossfn distillationlossfn def trainstepself data unpack data x y data forward pass of teacher teacherpredictions self teacherxdense teacherpredictions tf nn softmaxteacherpredictions axis1 with tf gradienttape as tape forward pass of student clspredictions distpredictions self student x 255 0 trainingtrue compute losses studentloss self studentlossfny clspredictions distillationloss self distillationlossfn teacherpredictions distpredictions loss studentloss distillationloss 2 compute gradients trainablevars self student trainablevariables gradients tape gradientloss trainablevars update weights self optimizer applygradientszipgradients trainablevars update the metrics configured in compile studentpredictions clspredictions distpredictions 2 self studentlosstracker updatestatestudentloss self distillationlosstracker updatestatedistillationloss self accuracy updatestatey studentpredictions return a dict of performance return m name m result for m in self metrics def teststepself data unpack the data x y data compute predictions yprediction self studentx 255 0 calculate the loss studentloss self studentlossfny yprediction update the metrics self studentlosstracker updatestatestudentloss self accuracy updatestatey yprediction return a dict of performance results m name m result for m in self metrics return results def callself inputs return self studentinputs 255 0 load the teacher model bitteacherflowers keras layers tfsmlayer filepathbitteacherflowers callendpointservingdefault training through distillation deittiny vitdistilled deitdistiller deitstudentdeittiny teacherbitteacherflowers lrscaled baselr 512 batchsize deitdistiller compile optimizerkeras optimizers adamw weightdecayweightdecay learningratelrscaled studentlossfnkeras losses categoricalcrossentropy fromlogitstrue labelsmoothing0 1 distillationlossfnkeras losses categoricalcrossentropyfromlogitstrue deitdistiller fit traindataset validationdatavaldataset epochsnumepochs notes acknowledgements introduction in the original vision transformers vit paper dosovitskiy et al https arxiv org abs 2010 11929 the s concluded that to perform on par with convolutional neural networks cnns vits need to be pre trained on larger datasets the larger the better this is mainly due to the lack of inductive biases in the vit architecture unlike cnns they don t have layers that exploit locality in a follow up paper steiner et al https arxiv org abs 2106 10270 the s show that it is possible to substantially improve the performance of vits with stronger regularization and longer training many groups have proposed different ways to deal with the problem of data intensiveness of vit training one such way was shown in the data efficient image transformers deit paper touvron et al https arxiv org abs 2012 12877 the s introduced a distillation technique that is specific to transformer based vision models deit is among the first works to show that it s possible to train vits well without using larger datasets in this example we implement the distillation recipe proposed in deit this requires us to slightly tweak the original vit architecture and write a custom training loop to implement the distillation recipe to comfortably navigate through this example you ll be expected to know how a vit and knowledge distillation work the following are good resources in case you needed a refresher vit on keras io https keras io examples vision image_classification_with_vision_transformer knowledge distillation on keras io https keras io examples vision knowledge_distillation imports constants model training data you probably noticed that dropout_rate has been set 0 0 dropout has been used in the implementation to keep it complete for smaller models like the one used in this example you don t need it but for bigger models using dropout helps load the tf_flowers dataset and prepare preprocessing utilities the s use an array of different augmentation techniques including mixup zhang et al https arxiv org abs 1710 09412 randaugment cubuk et al https arxiv org abs 1909 13719 and so on however to keep the example simple to work through we ll discard them resize to a bigger spatial resolution and take the random crops implementing the deit variants of vit since deit is an extension of vit it d make sense to first implement vit and then extend it to support deit s components first we ll implement a layer for stochastic depth huang et al https arxiv org abs 1603 09382 which is used in deit for regularization referred from github com rwightman pytorch image models now we ll implement the mlp and transformer blocks ffn for a transformer block iterate over the hidden units and add dense dropout transformer block with pre norm layer normalization 1 multi head self attention layer 1 skip connection 1 layer normalization 2 mlp layer 1 skip connection 2 we ll now implement a vitclassifier class building on top of the components we just developed here we ll be following the original pooling strategy used in the vit paper use a class token and use the feature representations corresponding to it for classification vision transformer base class patchify linear projection reshaping positional embedding transformer blocks cls token other layers create patches and project the patches append class token if needed add positional embeddings to the projected patches b number_patches projection_dim iterate over the number of layers and stack up blocks of transformer add a transformer block final layer normalization pool representation classification head this class can be used standalone as vit and is end to end trainable just remove the distilled phrase in model_type and it should work with vit_tiny vitclassifier let s now extend it to deit the following figure presents the schematic of deit taken from the deit paper https i imgur com 5lmg2xs png apart from the class token deit has another token for distillation during distillation the logits corresponding to the class token are compared to the true labels and the logits corresponding to the distillation token are compared to the teacher s predictions cls and distillation tokens positional embedding head layers create patches and project the patches append the tokens add positional embeddings to the projected patches b number_patches projection_dim iterate over the number of layers and stack up blocks of transformer add a transformer block final layer normalization classification heads during standard train finetune inference average the classifier predictions only return separate classification predictions when training in distilled mode let s verify if the vitdistilled class can be initialized and called as expected implementing the trainer unlike what happens in standard knowledge distillation hinton et al https arxiv org abs 1503 02531 where a temperature scaled softmax is used as well as kl divergence deit s use the following loss function https i imgur com bxdxsbq png here ce is cross entropy psi is the softmax function z_s denotes student predictions y denotes true labels y_t denotes teacher predictions reference https keras io examples vision knowledge_distillation unpack data forward pass of teacher forward pass of student compute losses compute gradients update weights update the metrics configured in compile return a dict of performance unpack the data compute predictions calculate the loss update the metrics return a dict of performance load the teacher model this model is based on the bit family of resnets kolesnikov et al https arxiv org abs 1912 11370 fine tuned on the tf_flowers dataset you can refer to this notebook https github com sayakpaul deit tf blob main notebooks bit teacher ipynb to know how the training was performed the teacher model has about 212 million parameters which is about 40x more than the student shell wget q https github com sayakpaul deit tf releases download v0 1 0 bit_teacher_flowers zip unzip q bit_teacher_flowers zip training through distillation if we had trained the same model the vitclassifier from scratch with the exact same hyperparameters the model would have scored about 59 accuracy you can adapt the following code to reproduce this result vit_tiny vitclassifier inputs keras input resolution resolution 3 x keras layers rescaling scale 1 255 inputs outputs deit_tiny x model keras model inputs outputs model compile model fit notes through the use of distillation we re effectively transferring the inductive biases of a cnn based teacher model interestingly enough this distillation strategy works better with a cnn as the teacher model rather than a transformer as shown in the paper the use of regularization to train deit models is very important vit models are initialized with a combination of different initializers including truncated normal random normal glorot uniform etc if you re looking for end to end reproduction of the original results don t forget to initialize the vits well if you want to explore the pre trained deit models in tensorflow and keras with code for fine tuning check out these models on tf hub https tfhub dev sayakpaul collections deit 1 acknowledgements ross wightman for keeping timm https github com rwightman pytorch image models updated with readable implementations i referred to the implementations of vit and deit a lot during implementing them in tensorflow aritra roy gosthipaty https github com arig23498 who implemented some portions of the vitclassifier in another project google developers experts https developers google com programs experts program for supporting me with gcp credits which were used to run experiments for this example
import os

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

import tensorflow as tf
import tensorflow_datasets as tfds
import keras
from keras import layers

tfds.disable_progress_bar()
keras.utils.set_random_seed(42)




MODEL_TYPE = "deit_distilled_tiny_patch16_224"
RESOLUTION = 224
PATCH_SIZE = 16
NUM_PATCHES = (RESOLUTION // PATCH_SIZE) ** 2
LAYER_NORM_EPS = 1e-6
PROJECTION_DIM = 192
NUM_HEADS = 3
NUM_LAYERS = 12
MLP_UNITS = [
    PROJECTION_DIM * 4,
    PROJECTION_DIM,
]
DROPOUT_RATE = 0.0
DROP_PATH_RATE = 0.1


NUM_EPOCHS = 20
BASE_LR = 0.0005
WEIGHT_DECAY = 0.0001


BATCH_SIZE = 256
AUTO = tf.data.AUTOTUNE
NUM_CLASSES = 5






def preprocess_dataset(is_training=True):
    def fn(image, label):
        if is_training:
            
            
            image = tf.image.resize(image, (RESOLUTION + 20, RESOLUTION + 20))
            image = tf.image.random_crop(image, (RESOLUTION, RESOLUTION, 3))
            image = tf.image.random_flip_left_right(image)
        else:
            image = tf.image.resize(image, (RESOLUTION, RESOLUTION))
        label = tf.one_hot(label, depth=NUM_CLASSES)
        return image, label

    return fn


def prepare_dataset(dataset, is_training=True):
    if is_training:
        dataset = dataset.shuffle(BATCH_SIZE * 10)
    dataset = dataset.map(
        preprocess_dataset(is_training), num_parallel_calls=AUTO
    )
    return dataset.batch(BATCH_SIZE).prefetch(AUTO)


train_dataset, val_dataset = tfds.load(
    "tf_flowers", split=["train[:90%]", "train[90%:]"], as_supervised=True
)
num_train = train_dataset.cardinality()
num_val = val_dataset.cardinality()
print(f"Number of training examples: {num_train}")
print(f"Number of validation examples: {num_val}")

train_dataset = prepare_dataset(train_dataset, is_training=True)
val_dataset = prepare_dataset(val_dataset, is_training=False)





class StochasticDepth(layers.Layer):
    def __init__(self, drop_prop, **kwargs):
        super().__init__(**kwargs)
        self.drop_prob = drop_prop

    def call(self, x, training=True):
        if training:
            keep_prob = 1 - self.drop_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





def mlp(x, dropout_rate: float, hidden_units):
    
    
    
    for idx, units in enumerate(hidden_units):
        x = layers.Dense(
            units,
            activation=tf.nn.gelu if idx == 0 else None,
        )(x)
        x = layers.Dropout(dropout_rate)(x)
    return x


def transformer(drop_prob: float, name: str) -> keras.Model:
    
    num_patches = (
        NUM_PATCHES + 2 if "distilled" in MODEL_TYPE else NUM_PATCHES + 1
    )
    encoded_patches = layers.Input((num_patches, PROJECTION_DIM))

    
    x1 = layers.LayerNormalization(epsilon=LAYER_NORM_EPS)(encoded_patches)

    
    attention_output = layers.MultiHeadAttention(
        num_heads=NUM_HEADS,
        key_dim=PROJECTION_DIM,
        dropout=DROPOUT_RATE,
    )(x1, x1)
    attention_output = (
        StochasticDepth(drop_prob)(attention_output)
        if drop_prob
        else attention_output
    )

    
    x2 = layers.Add()([attention_output, encoded_patches])

    
    x3 = layers.LayerNormalization(epsilon=LAYER_NORM_EPS)(x2)

    
    x4 = mlp(x3, hidden_units=MLP_UNITS, dropout_rate=DROPOUT_RATE)
    x4 = StochasticDepth(drop_prob)(x4) if drop_prob else x4

    
    outputs = layers.Add()([x2, x4])

    return keras.Model(encoded_patches, outputs, name=name)





class ViTClassifier(keras.Model):
    

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

        
        self.projection = keras.Sequential(
            [
                layers.Conv2D(
                    filters=PROJECTION_DIM,
                    kernel_size=(PATCH_SIZE, PATCH_SIZE),
                    strides=(PATCH_SIZE, PATCH_SIZE),
                    padding="VALID",
                    name="conv_projection",
                ),
                layers.Reshape(
                    target_shape=(NUM_PATCHES, PROJECTION_DIM),
                    name="flatten_projection",
                ),
            ],
            name="projection",
        )

        
        init_shape = (
            1,
            NUM_PATCHES + 1,
            PROJECTION_DIM,
        )
        self.positional_embedding = tf.Variable(
            tf.zeros(init_shape), name="pos_embedding"
        )

        
        dpr = [x for x in tf.linspace(0.0, DROP_PATH_RATE, NUM_LAYERS)]
        self.transformer_blocks = [
            transformer(drop_prob=dpr[i], name=f"transformer_block_{i}")
            for i in range(NUM_LAYERS)
        ]

        
        initial_value = tf.zeros((1, 1, PROJECTION_DIM))
        self.cls_token = tf.Variable(
            initial_value=initial_value, trainable=True, name="cls"
        )

        
        self.dropout = layers.Dropout(DROPOUT_RATE)
        self.layer_norm = layers.LayerNormalization(epsilon=LAYER_NORM_EPS)
        self.head = layers.Dense(
            NUM_CLASSES,
            name="classification_head",
        )

    def call(self, inputs):
        n = tf.shape(inputs)[0]

        
        projected_patches = self.projection(inputs)

        
        cls_token = tf.tile(self.cls_token, (n, 1, 1))
        cls_token = tf.cast(cls_token, projected_patches.dtype)
        projected_patches = tf.concat([cls_token, projected_patches], axis=1)

        
        encoded_patches = (
            self.positional_embedding + projected_patches
        )  
        encoded_patches = self.dropout(encoded_patches)

        
        
        for transformer_module in self.transformer_blocks:
            
            encoded_patches = transformer_module(encoded_patches)

        
        representation = self.layer_norm(encoded_patches)

        
        encoded_patches = representation[:, 0]

        
        output = self.head(encoded_patches)
        return output





class ViTDistilled(ViTClassifier):
    def __init__(self, regular_training=False, **kwargs):
        super().__init__(**kwargs)
        self.num_tokens = 2
        self.regular_training = regular_training

        
        init_value = tf.zeros((1, 1, PROJECTION_DIM))
        self.dist_token = tf.Variable(init_value, name="dist_token")
        self.positional_embedding = tf.Variable(
            tf.zeros(
                (
                    1,
                    NUM_PATCHES + self.num_tokens,
                    PROJECTION_DIM,
                )
            ),
            name="pos_embedding",
        )

        
        self.head = layers.Dense(
            NUM_CLASSES,
            name="classification_head",
        )
        self.head_dist = layers.Dense(
            NUM_CLASSES,
            name="distillation_head",
        )

    def call(self, inputs, training=False):
        n = tf.shape(inputs)[0]

        
        projected_patches = self.projection(inputs)

        
        cls_token = tf.tile(self.cls_token, (n, 1, 1))
        dist_token = tf.tile(self.dist_token, (n, 1, 1))
        cls_token = tf.cast(cls_token, projected_patches.dtype)
        dist_token = tf.cast(dist_token, projected_patches.dtype)
        projected_patches = tf.concat(
            [cls_token, dist_token, projected_patches], axis=1
        )

        
        encoded_patches = (
            self.positional_embedding + projected_patches
        )  
        encoded_patches = self.dropout(encoded_patches)

        
        
        for transformer_module in self.transformer_blocks:
            
            encoded_patches = transformer_module(encoded_patches)

        
        representation = self.layer_norm(encoded_patches)

        
        x, x_dist = (
            self.head(representation[:, 0]),
            self.head_dist(representation[:, 1]),
        )

        if not training or self.regular_training:
            
            
            return (x + x_dist) / 2

        elif training:
            
            
            return x, x_dist




deit_tiny_distilled = ViTDistilled()

dummy_inputs = tf.ones((2, 224, 224, 3))
outputs = deit_tiny_distilled(dummy_inputs, training=False)
print(f"output_shape: {outputs.shape}")




class DeiT(keras.Model):
    
    
    def __init__(self, student, teacher, **kwargs):
        super().__init__(**kwargs)
        self.student = student
        self.teacher = teacher
        self.student_loss_tracker = keras.metrics.Mean(name="student_loss")
        self.distillation_loss_tracker = keras.metrics.Mean(
            name="distillation_loss"
        )
        self.accuracy = keras.metrics.CategoricalAccuracy(name="accuracy")

    @property
    def metrics(self):
        return [
            self.accuracy,
            self.student_loss_tracker,
            self.distillation_loss_tracker,
        ]

    def compile(
        self,
        optimizer,
        student_loss_fn,
        distillation_loss_fn,
        run_eagerly=False,
        jit_compile=False,
    ):
        super().compile(
            optimizer=optimizer,
            run_eagerly=run_eagerly,
            jit_compile=jit_compile,
        )
        self.student_loss_fn = student_loss_fn
        self.distillation_loss_fn = distillation_loss_fn

    def train_step(self, data):
        
        x, y = data

        
        teacher_predictions = self.teacher(x)["dense"]
        teacher_predictions = tf.nn.softmax(teacher_predictions, axis=-1)

        with tf.GradientTape() as tape:
            
            cls_predictions, dist_predictions = self.student(
                x / 255.0, training=True
            )

            
            student_loss = self.student_loss_fn(y, cls_predictions)
            distillation_loss = self.distillation_loss_fn(
                teacher_predictions, dist_predictions
            )
            loss = (student_loss + distillation_loss) / 2

        
        trainable_vars = self.student.trainable_variables
        gradients = tape.gradient(loss, trainable_vars)

        
        self.optimizer.apply_gradients(zip(gradients, trainable_vars))

        
        student_predictions = (cls_predictions + dist_predictions) / 2
        self.student_loss_tracker.update_state(student_loss)
        self.distillation_loss_tracker.update_state(distillation_loss)
        self.accuracy.update_state(y, student_predictions)

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

    def test_step(self, data):
        
        x, y = data

        
        y_prediction = self.student(x / 255.0)

        
        student_loss = self.student_loss_fn(y, y_prediction)

        
        self.student_loss_tracker.update_state(student_loss)
        self.accuracy.update_state(y, y_prediction)

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

    def call(self, inputs):
        return self.student(inputs / 255.0)






bit_teacher_flowers = keras.layers.TFSMLayer(
    filepath="bit_teacher_flowers",
    call_endpoint="serving_default",
)




deit_tiny = ViTDistilled()
deit_distiller = DeiT(student=deit_tiny, teacher=bit_teacher_flowers)

lr_scaled = (BASE_LR / 512) * BATCH_SIZE
deit_distiller.compile(
    optimizer=keras.optimizers.AdamW(
        weight_decay=WEIGHT_DECAY, learning_rate=lr_scaled
    ),
    student_loss_fn=keras.losses.CategoricalCrossentropy(
        from_logits=True, label_smoothing=0.1
    ),
    distillation_loss_fn=keras.losses.CategoricalCrossentropy(from_logits=True),
)
_ = deit_distiller.fit(
    train_dataset, validation_data=val_dataset, epochs=NUM_EPOCHS
)