title knowledge distillation kenneth boruphttps twitter comkennethborup converted to keras 3 by md awsafur rahmanhttps awsaf49 github io date created 20200901 last modified 20200901 description implementation of classical knowledge distillation accelerator gpu introduction to knowledge distillation knowledge distillation is a procedure for model compression in which a small student model is trained to match a large pretrained teacher model knowledge is transferred from the teacher model to the student by minimizing a loss function aimed at matching softened teacher logits as well as groundtruth labels the logits are softened by applying a temperature scaling function in the softmax effectively smoothing out the probability distribution and revealing interclass relationships learned by the teacher reference hinton et al 2015https arxiv orgabs1503 02531 setup construct distiller class the custom distiller class overrides the model methods compile computeloss and call in order to use the distiller we need a trained teacher model a student model to train a student loss function on the difference between student predictions and groundtruth a distillation loss function along with a temperature on the difference between the soft student predictions and the soft teacher labels an alpha factor to weight the student and distillation loss an optimizer for the student and optional metrics to evaluate performance in the computeloss method we perform a forward pass of both the teacher and student calculate the loss with weighting of the studentloss and distillationloss by alpha and 1 alpha respectively note only the student weights are updated configure the distiller args optimizer keras optimizer for the student weights metrics keras metrics for evaluation studentlossfn loss function of difference between student predictions and groundtruth distillationlossfn loss function of difference between soft student predictions and soft teacher predictions alpha weight to studentlossfn and 1alpha to distillationlossfn temperature temperature for softening probability distributions larger temperature gives softer distributions create student and teacher models initialy we create a teacher model and a smaller student model both models are convolutional neural networks and created using sequential but could be any keras model create the teacher create the student clone student for later comparison prepare the dataset the dataset used for training the teacher and distilling the teacher is mnisthttps keras ioapidatasetsmnist and the procedure would be equivalent for any other dataset e g cifar10https keras ioapidatasetscifar10 with a suitable choice of models both the student and teacher are trained on the training set and evaluated on the test set prepare the train and test dataset normalize data train the teacher in knowledge distillation we assume that the teacher is trained and fixed thus we start by training the teacher model on the training set in the usual way train teacher as usual train and evaluate teacher on data distill teacher to student we have already trained the teacher model and we only need to initialize a distillerstudent teacher instance compile it with the desired losses hyperparameters and optimizer and distill the teacher to the student initialize and compile distiller distill teacher to student evaluate student on test dataset train student from scratch for comparison we can also train an equivalent student model from scratch without the teacher in order to evaluate the performance gain obtained by knowledge distillation train student as doen usually train and evaluate student trained from scratch if the teacher is trained for 5 full epochs and the student is distilled on this teacher for 3 full epochs you should in this example experience a performance boost compared to training the same student model from scratch and even compared to the teacher itself you should expect the teacher to have accuracy around 97 6 the student trained from scratch should be around 97 6 and the distilled student should be around 98 1 remove or try out different seeds to use different weight initializations introduction to knowledge distillation knowledge distillation is a procedure for model compression in which a small student model is trained to match a large pre trained teacher model knowledge is transferred from the teacher model to the student by minimizing a loss function aimed at matching softened teacher logits as well as ground truth labels the logits are softened by applying a temperature scaling function in the softmax effectively smoothing out the probability distribution and revealing inter class relationships learned by the teacher reference hinton et al 2015 https arxiv org abs 1503 02531 setup construct distiller class the custom distiller class overrides the model methods compile compute_loss and call in order to use the distiller we need a trained teacher model a student model to train a student loss function on the difference between student predictions and ground truth a distillation loss function along with a temperature on the difference between the soft student predictions and the soft teacher labels an alpha factor to weight the student and distillation loss an optimizer for the student and optional metrics to evaluate performance in the compute_loss method we perform a forward pass of both the teacher and student calculate the loss with weighting of the student_loss and distillation_loss by alpha and 1 alpha respectively note only the student weights are updated configure the distiller args optimizer keras optimizer for the student weights metrics keras metrics for evaluation student_loss_fn loss function of difference between student predictions and ground truth distillation_loss_fn loss function of difference between soft student predictions and soft teacher predictions alpha weight to student_loss_fn and 1 alpha to distillation_loss_fn temperature temperature for softening probability distributions larger temperature gives softer distributions create student and teacher models initialy we create a teacher model and a smaller student model both models are convolutional neural networks and created using sequential but could be any keras model create the teacher create the student clone student for later comparison prepare the dataset the dataset used for training the teacher and distilling the teacher is mnist https keras io api datasets mnist and the procedure would be equivalent for any other dataset e g cifar 10 https keras io api datasets cifar10 with a suitable choice of models both the student and teacher are trained on the training set and evaluated on the test set prepare the train and test dataset normalize data train the teacher in knowledge distillation we assume that the teacher is trained and fixed thus we start by training the teacher model on the training set in the usual way train teacher as usual train and evaluate teacher on data distill teacher to student we have already trained the teacher model and we only need to initialize a distiller student teacher instance compile it with the desired losses hyperparameters and optimizer and distill the teacher to the student initialize and compile distiller distill teacher to student evaluate student on test dataset train student from scratch for comparison we can also train an equivalent student model from scratch without the teacher in order to evaluate the performance gain obtained by knowledge distillation train student as doen usually train and evaluate student trained from scratch if the teacher is trained for 5 full epochs and the student is distilled on this teacher for 3 full epochs you should in this example experience a performance boost compared to training the same student model from scratch and even compared to the teacher itself you should expect the teacher to have accuracy around 97 6 the student trained from scratch should be around 97 6 and the distilled student should be around 98 1 remove or try out different seeds to use different weight initializations
import os

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

import keras
from keras import layers
from keras import ops
import tensorflow as tf
import numpy as np




class Distiller(keras.Model):
    def __init__(self, student, teacher):
        super().__init__()
        self.teacher = teacher
        self.student = student

    def compile(
        self,
        optimizer,
        metrics,
        student_loss_fn,
        distillation_loss_fn,
        alpha=0.1,
        temperature=3,
    ):
        
        super().compile(optimizer=optimizer, metrics=metrics)
        self.student_loss_fn = student_loss_fn
        self.distillation_loss_fn = distillation_loss_fn
        self.alpha = alpha
        self.temperature = temperature

    def compute_loss(
        self, x=None, y=None, y_pred=None, sample_weight=None, allow_empty=False
    ):
        teacher_pred = self.teacher(x, training=False)
        student_loss = self.student_loss_fn(y, y_pred)

        distillation_loss = self.distillation_loss_fn(
            ops.softmax(teacher_pred / self.temperature, axis=1),
            ops.softmax(y_pred / self.temperature, axis=1),
        ) * (self.temperature**2)

        loss = self.alpha * student_loss + (1 - self.alpha) * distillation_loss
        return loss

    def call(self, x):
        return self.student(x)





teacher = keras.Sequential(
    [
        keras.Input(shape=(28, 28, 1)),
        layers.Conv2D(256, (3, 3), strides=(2, 2), padding="same"),
        layers.LeakyReLU(negative_slope=0.2),
        layers.MaxPooling2D(pool_size=(2, 2), strides=(1, 1), padding="same"),
        layers.Conv2D(512, (3, 3), strides=(2, 2), padding="same"),
        layers.Flatten(),
        layers.Dense(10),
    ],
    name="teacher",
)


student = keras.Sequential(
    [
        keras.Input(shape=(28, 28, 1)),
        layers.Conv2D(16, (3, 3), strides=(2, 2), padding="same"),
        layers.LeakyReLU(negative_slope=0.2),
        layers.MaxPooling2D(pool_size=(2, 2), strides=(1, 1), padding="same"),
        layers.Conv2D(32, (3, 3), strides=(2, 2), padding="same"),
        layers.Flatten(),
        layers.Dense(10),
    ],
    name="student",
)


student_scratch = keras.models.clone_model(student)




batch_size = 64
(x_train, y_train), (x_test, y_test) = keras.datasets.mnist.load_data()


x_train = x_train.astype("float32") / 255.0
x_train = np.reshape(x_train, (-1, 28, 28, 1))

x_test = x_test.astype("float32") / 255.0
x_test = np.reshape(x_test, (-1, 28, 28, 1))





teacher.compile(
    optimizer=keras.optimizers.Adam(),
    loss=keras.losses.SparseCategoricalCrossentropy(from_logits=True),
    metrics=[keras.metrics.SparseCategoricalAccuracy()],
)


teacher.fit(x_train, y_train, epochs=5)
teacher.evaluate(x_test, y_test)




distiller = Distiller(student=student, teacher=teacher)
distiller.compile(
    optimizer=keras.optimizers.Adam(),
    metrics=[keras.metrics.SparseCategoricalAccuracy()],
    student_loss_fn=keras.losses.SparseCategoricalCrossentropy(
        from_logits=True
    ),
    distillation_loss_fn=keras.losses.KLDivergence(),
    alpha=0.1,
    temperature=10,
)


distiller.fit(x_train, y_train, epochs=3)


distiller.evaluate(x_test, y_test)




student_scratch.compile(
    optimizer=keras.optimizers.Adam(),
    loss=keras.losses.SparseCategoricalCrossentropy(from_logits=True),
    metrics=[keras.metrics.SparseCategoricalAccuracy()],
)


student_scratch.fit(x_train, y_train, epochs=3)
student_scratch.evaluate(x_test, y_test)