Adds Poisson, KL Divergence Probabilistic Metric / Loss (keras-team#62)

* Adds Poisson, KL Divergence Probabilistic Metric / Loss

* Update format and docstring

Author: sampathweb

keras_core/losses/__init__.py

@@ -2,36 +2,44 @@
 from keras_core.losses.loss import Loss
 from keras_core.losses.losses import CategoricalHinge
 from keras_core.losses.losses import Hinge
+from keras_core.losses.losses import KLDivergence
 from keras_core.losses.losses import LossFunctionWrapper
 from keras_core.losses.losses import MeanAbsoluteError
 from keras_core.losses.losses import MeanAbsolutePercentageError
 from keras_core.losses.losses import MeanSquaredError
 from keras_core.losses.losses import MeanSquaredLogarithmicError
+from keras_core.losses.losses import Poisson
 from keras_core.losses.losses import SquaredHinge
 from keras_core.losses.losses import categorical_hinge
 from keras_core.losses.losses import hinge
+from keras_core.losses.losses import kl_divergence
 from keras_core.losses.losses import mean_absolute_error
 from keras_core.losses.losses import mean_absolute_percentage_error
 from keras_core.losses.losses import mean_squared_error
 from keras_core.losses.losses import mean_squared_logarithmic_error
+from keras_core.losses.losses import poisson
 from keras_core.losses.losses import squared_hinge
 from keras_core.saving import serialization_lib
 
 ALL_OBJECTS = {
     Loss,
     LossFunctionWrapper,
+    KLDivergence,
     MeanSquaredError,
     MeanAbsoluteError,
     MeanAbsolutePercentageError,
     MeanSquaredLogarithmicError,
     Hinge,
+    Poisson,
     SquaredHinge,
     CategoricalHinge,
+    kl_divergence,
     mean_squared_error,
     mean_absolute_error,
     mean_absolute_percentage_error,
     mean_squared_logarithmic_error,
     hinge,
+    poisson,
     squared_hinge,
     categorical_hinge,
 }

keras_core/losses/losses.py

@@ -264,6 +264,54 @@ def get_config(self):
         return Loss.get_config(self)
 
 
+@keras_core_export("keras_core.losses.KLDivergence")
+class KLDivergence(LossFunctionWrapper):
+    """Computes Kullback-Leibler divergence loss between `y_true` & `y_pred`.
+
+    Formula:
+
+    ```python
+    loss = y_true * log(y_true / y_pred)
+    ```
+
+    Args:
+        reduction: Type of reduction to apply to loss. For almost all cases
+            this defaults to `"sum_over_batch_size"`. Options are `"sum"`,
+            `"sum_over_batch_size"` or `None`.
+        name: Optional name for the instance. Defaults to 'kl_divergence'.
+    """
+
+    def __init__(self, reduction="sum_over_batch_size", name="kl_divergence"):
+        super().__init__(kl_divergence, reduction=reduction, name=name)
+
+    def get_config(self):
+        return Loss.get_config(self)
+
+
+@keras_core_export("keras_core.losses.Poisson")
+class Poisson(LossFunctionWrapper):
+    """Computes the Poisson loss between `y_true` & `y_pred`.
+
+    Formula:
+
+    ```python
+    loss = y_pred - y_true * log(y_pred)
+    ```
+
+    Args:
+        reduction: Type of reduction to apply to loss. For almost all cases
+            this defaults to `"sum_over_batch_size"`. Options are `"sum"`,
+            `"sum_over_batch_size"` or `None`.
+        name: Optional name for the instance. Defaults to `"poisson"`
+    """
+
+    def __init__(self, reduction="sum_over_batch_size", name="poisson"):
+        super().__init__(poisson, reduction=reduction, name=name)
+
+    def get_config(self):
+        return Loss.get_config(self)
+
+
 def convert_binary_labels_to_hinge(y_true):
     """Converts binary labels into -1/1 for hinge loss/metric calculation."""
     are_zeros = ops.equal(y_true, 0)
@@ -582,3 +630,82 @@ def cosine_similarity(y_true, y_pred, axis=-1):
     y_pred = normalize(y_pred, axis=axis)
     y_true = normalize(y_true, axis=axis)
     return -ops.sum(y_true * y_pred, axis=axis)
+
+
+@keras_core_export(
+    [
+        "keras_core.metrics.kl_divergence",
+        "keras_core.losses.kl_divergence",
+    ]
+)
+def kl_divergence(y_true, y_pred):
+    """Computes Kullback-Leibler divergence loss between `y_true` & `y_pred`.
+
+    Formula:
+
+    ```python
+    loss = y_true * log(y_true / y_pred)
+    ```
+
+    Standalone usage:
+
+    >>> y_true = np.random.randint(0, 2, size=(2, 3)).astype(np.float32)
+    >>> y_pred = np.random.random(size=(2, 3))
+    >>> loss = keras_core.losses.kl_divergence(y_true, y_pred)
+    >>> assert loss.shape == (2,)
+    >>> y_true = ops.clip(y_true, 1e-7, 1)
+    >>> y_pred = ops.clip(y_pred, 1e-7, 1)
+    >>> assert np.array_equal(
+    ...     loss, np.sum(y_true * np.log(y_true / y_pred), axis=-1))
+
+    Args:
+        y_true: Tensor of true targets.
+        y_pred: Tensor of predicted targets.
+
+    Returns:
+        KL Divergence loss values with shape = `[batch_size, d0, .. dN-1]`.
+    """
+    y_pred = ops.convert_to_tensor(y_pred)
+    y_true = ops.convert_to_tensor(y_true, y_pred.dtype)
+    y_true = ops.clip(y_true, backend.epsilon(), 1)
+    y_pred = ops.clip(y_pred, backend.epsilon(), 1)
+    return ops.sum(y_true * ops.log(y_true / y_pred), axis=-1)
+
+
+@keras_core_export(
+    [
+        "keras_core.metrics.poisson",
+        "keras_core.losses.poisson",
+    ]
+)
+def poisson(y_true, y_pred):
+    """Computes the Poisson loss between y_true and y_pred.
+
+    Formula:
+
+    ```python
+    loss = y_pred - y_true * log(y_pred)
+    ```
+
+    Standalone usage:
+
+    >>> y_true = np.random.randint(0, 2, size=(2, 3))
+    >>> y_pred = np.random.random(size=(2, 3))
+    >>> loss = keras_core.losses.poisson(y_true, y_pred)
+    >>> assert loss.shape == (2,)
+    >>> y_pred = y_pred + 1e-7
+    >>> assert np.allclose(
+    ...     loss, np.mean(y_pred - y_true * np.log(y_pred), axis=-1),
+    ...     atol=1e-5)
+
+    Args:
+        y_true: Ground truth values. shape = `[batch_size, d0, .. dN]`.
+        y_pred: The predicted values. shape = `[batch_size, d0, .. dN]`.
+
+    Returns:
+        Poisson loss values with shape = `[batch_size, d0, .. dN-1]`.
+    """
+    y_pred = ops.convert_to_tensor(y_pred)
+    y_true = ops.convert_to_tensor(y_true, dtype=y_pred.dtype)
+    epsilon = ops.convert_to_tensor(backend.epsilon())
+    return ops.mean(y_pred - y_true * ops.log(y_pred + epsilon), axis=-1)

keras_core/losses/losses_test.py

@@ -461,3 +461,171 @@ def test_axis(self):
         loss = cosine_obj(self.y_true, self.y_pred)
         expected_loss = -np.mean(self.expected_loss)
         self.assertAlmostEqual(loss, expected_loss, 3)
+
+
+class KLDivergenceTest(testing.TestCase):
+    def setup(self):
+        self.y_pred = np.asarray(
+            [0.4, 0.9, 0.12, 0.36, 0.3, 0.4], dtype=np.float32
+        ).reshape((2, 3))
+        self.y_true = np.asarray(
+            [0.5, 0.8, 0.12, 0.7, 0.43, 0.8], dtype=np.float32
+        ).reshape((2, 3))
+
+        self.batch_size = 2
+        self.expected_losses = np.multiply(
+            self.y_true, np.log(self.y_true / self.y_pred)
+        )
+
+    def test_config(self):
+        k_obj = losses.KLDivergence(reduction="sum", name="kld")
+        self.assertEqual(k_obj.name, "kld")
+        self.assertEqual(k_obj.reduction, "sum")
+
+    def test_unweighted(self):
+        self.setup()
+        k_obj = losses.KLDivergence()
+
+        loss = k_obj(self.y_true, self.y_pred)
+        expected_loss = np.sum(self.expected_losses) / self.batch_size
+        self.assertAlmostEqual(loss, expected_loss, 3)
+
+    def test_scalar_weighted(self):
+        self.setup()
+        k_obj = losses.KLDivergence()
+        sample_weight = 2.3
+
+        loss = k_obj(self.y_true, self.y_pred, sample_weight=sample_weight)
+        expected_loss = (
+            sample_weight * np.sum(self.expected_losses) / self.batch_size
+        )
+        self.assertAlmostEqual(loss, expected_loss, 3)
+
+        # Verify we get the same output when the same input is given
+        loss_2 = k_obj(self.y_true, self.y_pred, sample_weight=sample_weight)
+        self.assertAlmostEqual(loss, loss_2, 3)
+
+    def test_sample_weighted(self):
+        self.setup()
+        k_obj = losses.KLDivergence()
+        sample_weight = np.asarray([1.2, 3.4], dtype=np.float32).reshape((2, 1))
+        loss = k_obj(self.y_true, self.y_pred, sample_weight=sample_weight)
+
+        expected_loss = np.multiply(
+            self.expected_losses,
+            np.asarray(
+                [1.2, 1.2, 1.2, 3.4, 3.4, 3.4], dtype=np.float32
+            ).reshape(2, 3),
+        )
+        expected_loss = np.sum(expected_loss) / self.batch_size
+        self.assertAlmostEqual(loss, expected_loss, 3)
+
+    def test_timestep_weighted(self):
+        self.setup()
+        k_obj = losses.KLDivergence()
+        y_true = self.y_true.reshape(2, 3, 1)
+        y_pred = self.y_pred.reshape(2, 3, 1)
+        sample_weight = np.asarray([3, 6, 5, 0, 4, 2]).reshape(2, 3)
+        expected_losses = np.sum(
+            np.multiply(y_true, np.log(y_true / y_pred)), axis=-1
+        )
+        loss = k_obj(y_true, y_pred, sample_weight=sample_weight)
+
+        num_timesteps = 3
+        expected_loss = np.sum(expected_losses * sample_weight) / (
+            self.batch_size * num_timesteps
+        )
+        self.assertAlmostEqual(loss, expected_loss, 3)
+
+    def test_zero_weighted(self):
+        self.setup()
+        k_obj = losses.KLDivergence()
+        loss = k_obj(self.y_true, self.y_pred, sample_weight=0)
+        self.assertAlmostEqual(loss, 0.0, 3)
+
+
+class PoissonTest(testing.TestCase):
+    def setup(self):
+        self.y_pred = np.asarray([1, 9, 2, 5, 2, 6], dtype=np.float32).reshape(
+            (2, 3)
+        )
+        self.y_true = np.asarray([4, 8, 12, 8, 1, 3], dtype=np.float32).reshape(
+            (2, 3)
+        )
+
+        self.batch_size = 6
+        self.expected_losses = self.y_pred - np.multiply(
+            self.y_true, np.log(self.y_pred)
+        )
+
+    def test_config(self):
+        poisson_obj = losses.Poisson(reduction="sum", name="poisson")
+        self.assertEqual(poisson_obj.name, "poisson")
+        self.assertEqual(poisson_obj.reduction, "sum")
+
+    def test_unweighted(self):
+        self.setup()
+        poisson_obj = losses.Poisson()
+
+        loss = poisson_obj(self.y_true, self.y_pred)
+        expected_loss = np.sum(self.expected_losses) / self.batch_size
+        self.assertAlmostEqual(loss, expected_loss, 3)
+
+    def test_scalar_weighted(self):
+        self.setup()
+        poisson_obj = losses.Poisson()
+        sample_weight = 2.3
+        loss = poisson_obj(
+            self.y_true, self.y_pred, sample_weight=sample_weight
+        )
+        expected_loss = (
+            sample_weight * np.sum(self.expected_losses) / self.batch_size
+        )
+        self.assertAlmostEqual(loss, expected_loss, 3)
+        self.assertAlmostEqual(loss, expected_loss, 3)
+
+        # Verify we get the same output when the same input is given
+        loss_2 = poisson_obj(
+            self.y_true, self.y_pred, sample_weight=sample_weight
+        )
+        self.assertAlmostEqual(loss, loss_2, 3)
+
+    def test_sample_weighted(self):
+        self.setup()
+        poisson_obj = losses.Poisson()
+
+        sample_weight = np.asarray([1.2, 3.4]).reshape((2, 1))
+        loss = poisson_obj(
+            self.y_true, self.y_pred, sample_weight=sample_weight
+        )
+
+        expected_loss = np.multiply(
+            self.expected_losses,
+            np.asarray([1.2, 1.2, 1.2, 3.4, 3.4, 3.4]).reshape((2, 3)),
+        )
+        expected_loss = np.sum(expected_loss) / self.batch_size
+        self.assertAlmostEqual(loss, expected_loss, 3)
+
+    def test_timestep_weighted(self):
+        self.setup()
+        poisson_obj = losses.Poisson()
+        y_true = self.y_true.reshape(2, 3, 1)
+        y_pred = self.y_pred.reshape(2, 3, 1)
+        sample_weight = np.asarray([3, 6, 5, 0, 4, 2]).reshape(2, 3, 1)
+        expected_losses = y_pred - np.multiply(y_true, np.log(y_pred))
+
+        loss = poisson_obj(
+            y_true,
+            y_pred,
+            sample_weight=np.asarray(sample_weight).reshape((2, 3)),
+        )
+        expected_loss = (
+            np.sum(expected_losses * sample_weight) / self.batch_size
+        )
+        self.assertAlmostEqual(loss, expected_loss, 3)
+
+    def test_zero_weighted(self):
+        self.setup()
+        poisson_obj = losses.Poisson()
+        loss = poisson_obj(self.y_true, self.y_pred, sample_weight=0)
+        self.assertAlmostEqual(loss, 0.0, 3)

keras_core/metrics/__init__.py

@@ -8,6 +8,8 @@
 from keras_core.metrics.hinge_metrics import Hinge
 from keras_core.metrics.hinge_metrics import SquaredHinge
 from keras_core.metrics.metric import Metric
+from keras_core.metrics.probabilistic_metrics import KLDivergence
+from keras_core.metrics.probabilistic_metrics import Poisson
 from keras_core.metrics.reduction_metrics import Mean
 from keras_core.metrics.reduction_metrics import MeanMetricWrapper
 from keras_core.metrics.reduction_metrics import Sum
@@ -26,6 +28,8 @@
     Hinge,
     SquaredHinge,
     CategoricalHinge,
+    KLDivergence,
+    Poisson,
 }
 ALL_OBJECTS_DICT = {cls.__name__: cls for cls in ALL_OBJECTS}
 ALL_OBJECTS_DICT.update(