From c9e7dc59303bb4305b7861d46922a503c7d7c5d1 Mon Sep 17 00:00:00 2001 From: AtomicVar Date: Thu, 18 Jan 2024 11:23:44 +0800 Subject: [PATCH] Add binary focal cross entropy loss function --- docs/src/python/nn/losses.rst | 1 + python/mlx/nn/losses.py | 88 +++++++++++++++++++++++++++++++++++ python/tests/test_losses.py | 60 ++++++++++++++++++++++++ 3 files changed, 149 insertions(+) diff --git a/docs/src/python/nn/losses.rst b/docs/src/python/nn/losses.rst index 6c4327eb83..cc1b164398 100644 --- a/docs/src/python/nn/losses.rst +++ b/docs/src/python/nn/losses.rst @@ -10,6 +10,7 @@ Loss Functions :template: nn-module-template.rst binary_cross_entropy + binary_focal_cross_entropy cosine_similarity_loss cross_entropy gaussian_nll_loss diff --git a/python/mlx/nn/losses.py b/python/mlx/nn/losses.py index 0f0021050e..903e6f04ab 100644 --- a/python/mlx/nn/losses.py +++ b/python/mlx/nn/losses.py @@ -121,6 +121,94 @@ def binary_cross_entropy( return _reduce(loss, reduction) +def binary_focal_cross_entropy( + inputs: mx.array, + targets: mx.array, + with_logits: bool = True, + class_balancing: bool = False, + alpha: float = 0.25, + gamma: float = 2.0, + reduction: Reduction = "mean", +) -> float: + r""" + Calculate the binary focal cross entropy loss between true labels and predicted probabilities. + + According to `Lin et al., 2018`_, it helps to apply a "focal factor" to down-weight easy + examples and focus more on hard examples. + + The loss is given by: + + .. math:: + \text{FL}(p,y)= + \begin{cases} + -\alpha(1-p)^\gamma\log(p), & \text{ if } y=1 \\ + -(1-\alpha)p^\gamma\log(1-p), & \text{otherwise.} + \end{cases} + + where :math:`p` is the predicted probability, :math:`y` is the target value, :math:`\alpha` is + the weight balancing factor for class 1, and :math:`\gamma` is the focusing parameter. + + When ``class_balancing`` is ``False`` and ``gamma`` is ``0``, this function is equivalent to the + binary cross-entropy loss. + + Args: + inputs (array): The predicted values. If ``with_logits`` is ``True``, then + ``inputs`` are unnormalized logits. Otherwise, ``inputs`` are probabilities. + targets (array): Binary target values in {0, 1}. + with_logits (bool, optional): Whether ``inputs`` are logits. Default: ``True``. + class_balancing (bool, optional): Whether to apply class balancing. Default: ``False``. + alpha (float, optional): A weight balancing factor for class 1. + Default: ``0.25`` (`Lin et al., 2018`_). The weight for class 0 is ``1.0 - alpha``. + gamma (float, optional): A focusing parameter used to compute the focal factor. + Default: ``2.0`` (`Lin et al., 2018`_). + reduction (str, optional): Specifies the reduction to apply to the output: + ``'none'`` | ``'mean'`` | ``'sum'``. Default: ``'mean'``. + + .. _Lin et al., 2018: https://arxiv.org/pdf/1708.02002.pdf + + Returns: + array: The computed binary focal cross entropy loss. + + Examples: + >>> import mlx.core as mx + >>> from mlx.nn.losses import binary_focal_cross_entropy + + >>> # inputs are logits, without class balancing + >>> logits = mx.array([[0.105361, 0.223144], [1.20397, 0.916291]]) + >>> targets = mx.array([[0, 1], [1, 0]]) + >>> binary_focal_cross_entropy(logits, targets, gamma=2.0) + array(0.244057, dtype=float32) + + >>> # inputs are probabilities, with class balancing + >>> probs = mx.array([[0.1, 0.1], [0.4, 0.4]]) + >>> targets = mx.array([[0, 1], [1, 0]]) + >>> binary_focal_cross_entropy(probs, targets, class_balancing=True, alpha=0.25, gamma=2.0) + array(0.114078, dtype=float32) + """ + if inputs.shape != targets.shape: + raise ValueError( + f"Inputs shape {inputs.shape} does not match targets shape {targets.shape}." + ) + + if with_logits: + inputs = mx.sigmoid(inputs) + + focusing_pos = (1 - inputs) ** gamma + focusing_neg = inputs**gamma + + loss_pos = -(focusing_pos * mx.log(inputs) * targets) + if class_balancing: + loss_pos *= alpha + + loss_neg = -(focusing_neg * mx.log(1 - inputs) * (1 - targets)) + if class_balancing: + loss_neg *= 1 - alpha + + loss = loss_pos + loss_neg + + return _reduce(loss, reduction) + + def l1_loss( predictions: mx.array, targets: mx.array, reduction: Reduction = "mean" ) -> mx.array: diff --git a/python/tests/test_losses.py b/python/tests/test_losses.py index a23f264542..165155c4ec 100644 --- a/python/tests/test_losses.py +++ b/python/tests/test_losses.py @@ -160,6 +160,66 @@ def _test_probs_as_inputs(): _test_logits_as_inputs() _test_probs_as_inputs() + def test_binary_focal_cross_entropy(self): + # Test ground truth values are calculated by tf.keras.losses.BinaryFocalCrossentropy + + alpha = 0.25 # class balancing factor + gamma = 2.0 # focal factor + + def _test_logits_as_inputs(): + logits = mx.array([-1.6, 0.5, 2.9, -1.8]) + targets = mx.array([0, 1, 0, 1]) + + # Test with reduction 'none', without class balancing + losses = nn.losses.binary_focal_cross_entropy( + logits, targets, gamma=gamma, reduction="none" + ) + expected = mx.array([0.00518928, 0.06757348, 2.653519, 1.438211]) + self.assertTrue(mx.allclose(losses, expected)) + + # Test with reduction 'none', with class balancing + losses_balanced = nn.losses.binary_focal_cross_entropy( + logits, + targets, + class_balancing=True, + alpha=alpha, + gamma=gamma, + reduction="none", + ) + expected_balanced = mx.array( + [0.00389196, 0.01689337, 1.9901392, 0.35955274] + ) + self.assertTrue(mx.allclose(losses_balanced, expected_balanced)) + + def _test_probs_as_inputs(): + probs = mx.array([0.5, 0.6, 0.7, 0.8]) + targets = mx.array([0, 1, 0, 1]) + + # Test with reduction 'none', without class balancing + losses = nn.losses.binary_focal_cross_entropy( + probs, targets, with_logits=False, gamma=gamma, reduction="none" + ) + expected = mx.array([0.17328674, 0.08173206, 0.5899465, 0.00892573]) + self.assertTrue(mx.allclose(losses, expected)) + + # Test with reduction 'none', with class balancing + losses_balanced = nn.losses.binary_focal_cross_entropy( + probs, + targets, + with_logits=False, + class_balancing=True, + alpha=alpha, + gamma=gamma, + reduction="none", + ) + expected_balanced = mx.array( + [0.12996505, 0.02043301, 0.44245988, 0.00223143] + ) + self.assertTrue(mx.allclose(losses_balanced, expected_balanced)) + + _test_logits_as_inputs() + _test_probs_as_inputs() + def test_l1_loss(self): predictions = mx.array([0.5, 0.2, 0.9, 0.0]) targets = mx.array([0.5, 0.2, 0.9, 0.0])