Add Distillation API to Keras (keras-team#21572)

* initial code dump

* clean up the implementation of the distillation api

* code reformat

* final clean up

* pre commit

* address gemini review comments

* address gemini review comments

* add a way to save trained student model

* disable tests in numpy and openvino backends

* pre commit

* address comments

* address comments

* run pre-commit

* update distiller and strategies

* code reformat

* clean up

* code reformat

* remove multi output distiilation

* clean up after merge

* address comments

* update file names

* subclass logits distillation loss from feature distillation loss

* update docstrings

* add validation for feature extraction setup

* Fix distiller API to accept single strategy or list of strategies

* minor fixes

* fix tests

* address reveiw comments

* oh

* address review comments

* address review comments

* address review comments

* code reformat

* address gemini review comments

* address gemini comments

* address review comments

Author: divyashreepathihalli

keras/api/__init__.py

@@ -12,6 +12,7 @@
 from keras import config as config
 from keras import constraints as constraints
 from keras import datasets as datasets
+from keras import distillation as distillation
 from keras import distribution as distribution
 from keras import dtype_policies as dtype_policies
 from keras import export as export

keras/api/_tf_keras/keras/__init__.py

@@ -10,6 +10,7 @@
 from keras import config as config
 from keras import constraints as constraints
 from keras import datasets as datasets
+from keras import distillation as distillation
 from keras import distribution as distribution
 from keras import dtype_policies as dtype_policies
 from keras import export as export

keras/api/_tf_keras/keras/distillation/__init__.py

@@ -0,0 +1,16 @@
+"""DO NOT EDIT.
+
+This file was autogenerated. Do not edit it by hand,
+since your modifications would be overwritten.
+"""
+
+from keras.src.distillation.distillation_loss import (
+    DistillationLoss as DistillationLoss,
+)
+from keras.src.distillation.distillation_loss import (
+    FeatureDistillation as FeatureDistillation,
+)
+from keras.src.distillation.distillation_loss import (
+    LogitsDistillation as LogitsDistillation,
+)
+from keras.src.distillation.distiller import Distiller as Distiller

keras/api/distillation/__init__.py

@@ -0,0 +1,16 @@
+"""DO NOT EDIT.
+
+This file was autogenerated. Do not edit it by hand,
+since your modifications would be overwritten.
+"""
+
+from keras.src.distillation.distillation_loss import (
+    DistillationLoss as DistillationLoss,
+)
+from keras.src.distillation.distillation_loss import (
+    FeatureDistillation as FeatureDistillation,
+)
+from keras.src.distillation.distillation_loss import (
+    LogitsDistillation as LogitsDistillation,
+)
+from keras.src.distillation.distiller import Distiller as Distiller

keras/src/distillation/__init__.py

@@ -0,0 +1 @@
+"""Distillation module for knowledge distillation in Keras."""