Source code for akida_models.distiller

#!/usr/bin/env python
# *****************************************************************************
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ******************************************************************************
"""
Tools for Knowledge Distillation training.

Originated from https://keras.io/examples/vision/knowledge_distillation/.

Reference Hinton et al. (2015) https://arxiv.org/abs/1503.02531
"""

import tensorflow as tf
from functools import partial
from tensorflow import GradientTape
from keras import Model
from keras.losses import KLDivergence, CategoricalCrossentropy


[docs]class Distiller(Model): """ The class that will be used to train the student model using the distillation knowledge method. Reference `Hinton et al. (2015) <https://arxiv.org/abs/1503.02531>`_. Args: student (keras.Model): the student model teacher (keras.Model): the well trained teacher model alpha (float, optional): weight to student_loss_fn and 1-alpha to distillation_loss_fn. Defaults to 0.1 """ def __init__(self, student, teacher, alpha=0.1): super().__init__() self.teacher = teacher self.student = student self.student_loss_fn = None self.distillation_loss_fn = None self.alpha = alpha @property def base_model(self): return self.student @property def layers(self): return self.base_model.layers def compile(self, optimizer, metrics, student_loss_fn, distillation_loss_fn): """ Configure the distiller. Args: optimizer (keras.optimizers.Optimizer): Keras optimizer for the student weights metrics (keras.metrics.Metric): Keras metrics for evaluation student_loss_fn (keras.losses.Loss): loss function of difference between student predictions and ground-truth distillation_loss_fn (keras.losses.Loss): loss function of difference between student predictions and teacher predictions """ super().compile(optimizer=optimizer, metrics=metrics) self.student_loss_fn = student_loss_fn self.distillation_loss_fn = distillation_loss_fn def train_step(self, data): # Unpack data x, y = data # Forward pass of teacher teacher_predictions = self.teacher(x, training=False) with GradientTape() as tape: # Forward pass of student student_predictions = self.student(x, training=True) # Compute losses student_loss = self.student_loss_fn(y, student_predictions) distillation_loss = self.distillation_loss_fn( teacher_predictions, student_predictions) loss = self.alpha * student_loss + (1 - self.alpha) * distillation_loss # Compute gradients trainable_vars = self.student.trainable_variables gradients = tape.gradient(loss, trainable_vars) # Update weights self.optimizer.apply_gradients(zip(gradients, trainable_vars)) # Update the metrics configured in `compile()`. self.compiled_metrics.update_state(y, student_predictions) # Return a dict of performance results = {m.name: m.result() for m in self.metrics} results.update({ "student_loss": student_loss, "distillation_loss": distillation_loss }) return results def test_step(self, data): # Unpack the data x, y = data # Compute predictions y_prediction = self.student(x, training=False) # Calculate the loss student_loss = self.student_loss_fn(y, y_prediction) # Update the metrics. self.compiled_metrics.update_state(y, y_prediction) # Return a dict of performance results = {m.name: m.result() for m in self.metrics} results.update({"student_loss": student_loss}) return results def save(self, *args, **kwargs): return self.base_model.save(*args, **kwargs) def save_weights(self, *args, **kwargs): return self.base_model.save_weights(*args, **kwargs) def load_weights(self, *args, **kwargs): return self.base_model.load_weights(*args, **kwargs)
[docs]class DeitDistiller(Distiller): """Distiller class to train the student model using the Knowledge Distillation (KD) method, found on https://arxiv.org/pdf/2012.12877.pdf The main difference with the classic KD is that the student has to produce two potential classification outputs. This type of training is based on the assumption that each output has sufficiently interacted with the whole model, therefore the main architecture can be trained through two different sources, as follows: >>> output, output_kd = student(input) >>> output_tc = teacher(input) >>> student_loss = student_loss_fn(y_true, output) >>> distillation_loss = distillation_loss_fn(output_tc, output_kd) This means we will expect to have different inputs for each loss, unlike classical KD, where the student's prediction is shared for both losses. However, given that each classifier has interacted with the student model, the gradient of each loss will contribute to the update of the model weights according to the alpha percentage. Args: student (keras.Model): the student model teacher (keras.Model): the well trained teacher model alpha (float, optional): weight to student_loss_fn and 1-alpha to distillation_loss_fn. Defaults to 0.1 temperature (float, optional): if ``distiller_type`` when compile is equal to 'soft', this value will be used as temperature parameter of KLDistillationLoss. Defaults to 1.0. """ def __init__(self, student, *args, temperature=1.0, **kwargs): assert len(student.outputs) == 2, "Student must be a multi-output model, with 2 outputs" # Append an output with the sum of heads y = tf.math.add_n(student.outputs) / 2 _student = Model(student.inputs, student.outputs + [y], name=student.name) super().__init__(_student, *args, **kwargs) self._student = student self.temperature = temperature @property def base_model(self): return self._student def compile(self, optimizer, metrics, student_loss_fn, distiller_type): """ Configure the distiller. Args: optimizer (keras.optimizers.Optimizer): Keras optimizer for the student weights metrics (keras.metrics.Metric): Keras metrics for evaluation student_loss_fn (keras.losses.Loss): loss function of difference between student predictions and ground-truth distiller_type (str): loss function type to define the difference between student predictions and teacher-truth, within ['soft', 'hard', 'none'] which will result in performing KLDistillationLoss, CategoricalCrossentropy or student_loss_fn only respectively. """ assert distiller_type in ['soft', 'hard', 'none'] def _loss_forward(y_true, y_pred, loss_fn, index=0, **kwargs): if isinstance(y_pred, (tuple, list)): y_pred = y_pred[index] return loss_fn(y_true, y_pred, **kwargs) def _compile_distillation_loss_fn(): if distiller_type == "soft": distillation_loss_fn = KLDistillationLoss(temperature=self.temperature) else: # Follow https://arxiv.org/pdf/2012.12877.pdf, this variant takes the # hard decision of the teacher as a true label. Therefore, we add the # prediction encoder, as well as a label smoothing equal to 0.1 y = tf.math.softmax(self.teacher.outputs[0], axis=-1) self.teacher = Model(self.teacher.inputs, y, name=self.teacher.name) distillation_loss_fn = CategoricalCrossentropy( from_logits=True, label_smoothing=0.1) return partial(_loss_forward, loss_fn=distillation_loss_fn, index=1) if distiller_type == "none" or self.teacher is None: # In this case, we just train the first output of student self.teacher = distillation_loss_fn = None self.student = Model(self.student.inputs, self.student.outputs[0], name=self.student.name) self.student.compile(optimizer, student_loss_fn, metrics) else: distillation_loss_fn = _compile_distillation_loss_fn() student_loss_fn = partial(_loss_forward, loss_fn=student_loss_fn) super().compile(optimizer, metrics, student_loss_fn, distillation_loss_fn) def _update_metrics(self, metrics): # Rename keys in the result dictionary for a more explicit display return {k.replace('output_1_', 'head_').replace('output_2_', 'dist_head_') .replace('output_3_', ''): v for k, v in metrics.items()} def train_step(self, data): if self.teacher is None: return self.student.train_step(data) return self._update_metrics(super().train_step(data)) def test_step(self, data): if self.teacher is None: return self.student.test_step(data) return self._update_metrics(super().test_step(data))
[docs]class KLDistillationLoss(KLDivergence): """ The `KLDistillationLoss` is a simple wrapper around the KLDivergence loss that accepts raw predictions instead of probability distributions. Before invoking the KLDivergence loss, it converts the inputs predictions to probabilities by dividing them by a constant 'temperature' and applies a softmax. Args: temperature (float): temperature for softening probability distributions. Larger temperature gives softer distributions. """ def __init__(self, temperature=3): self.temperature = temperature super().__init__() def __call__(self, y_true, y_pred, sample_weight=None): # Following https://github.com/facebookresearch/deit/blob/main/losses.py#L63 # The result of KLDivergence must be scaled scale_factor = tf.constant(self.temperature ** 2, dtype=tf.float32) return super().__call__( tf.nn.softmax(y_true / self.temperature, axis=1), tf.nn.softmax(y_pred / self.temperature, axis=1)) * scale_factor