Repository: deel-ai/deel-lip (language: python)

Issue:
Element-wise losses for multi-GPU/TPU strategy
Keras losses API requires that subclasses return element-wise values, i.e. each element of a batch should return a single value. Then, a reduction operation performs either an average over the batch (default case) or a sum over the batch (useful for multi-GPU or TPU trainings).

This PR proposes a refactoring of all provided losses to fit the Keras losses API. For hinge losses, this is straightforward. For KR losses, it is more complicated: by definition, KR losses compute average over the batch. We then propose two element-wise KR formulations for standard training or multi-GPU/TPU training:
* a standard element-wise KR formulation that accepts classic one-hot encoded labels,
* a KR formulation appropriate for multi-GPU/TPU, where labels must be specifically pre-processed.

Multiple API changes are introduced:
* the KR loss is now subclassed from tf.keras Loss. Use `KR()(y_true, y_pred)` instead of `KR(y_true, y_pred)`
* the negative_KR function is removed. Use instead `HKR(alpha=0)`
* a new utils function `process_labels_for_KR_loss()` is added for label processing when using a multi-GPU/TPU strategy
* a new argument `multi_gpu` is added in all KR-based losses (KR, HKR, MulticlassKR, MulticlassHKR). When training with multi-GPU/TPU, this flag must be set to True and the labels must be processed using `process_labels_for_KR_loss()`.



Discussion / hints:

Fix patch:
diff --git a/deel/lip/__init__.py b/deel/lip/__init__.py
index d2aa5513..109a2572 100644
--- a/deel/lip/__init__.py
+++ b/deel/lip/__init__.py
@@ -8,6 +8,7 @@
 from . import initializers
 from . import layers
 from . import losses
+from . import metrics
 from .model import Sequential, Model, vanillaModel
 from . import normalizers
 from . import utils
diff --git a/deel/lip/metrics.py b/deel/lip/metrics.py
new file mode 100644
index 00000000..b2a6e441
--- /dev/null
+++ b/deel/lip/metrics.py
@@ -0,0 +1,292 @@
+# Copyright IRT Antoine de Saint Exupéry et Université Paul Sabatier Toulouse III - All
+# rights reserved. DEEL is a research program operated by IVADO, IRT Saint Exupéry,
+# CRIAQ and ANITI - https://www.deel.ai/
+# =====================================================================================
+"""
+This module contains metrics applicable in provable robustness. See
+https://arxiv.org/abs/2006.06520 and https://arxiv.org/abs/2108.04062 for more
+information.
+"""
+import math
+import tensorflow as tf
+from tensorflow.keras.losses import Loss
+from tensorflow.keras.losses import Reduction
+from tensorflow.keras.utils import register_keras_serializable
+
+
+def _delta_multiclass(y_true, y_pred):
+    r"""
+    Compute the non-normalized provable robustness factor in multiclass setup (need
+    to be adjusted with
+    lipschitz constant.
+
+    .. math::
+            \Delta(x) = f_l(x) - \max_{i \neq l} f_i(x)
+
+    Args:
+        y_true: true labels must be in {1,0} or in {1,-1} (no label smoothing allowed)
+        y_pred: network predictions
+
+    Returns: non-normalized provable robustness factor
+
+    """
+    ynl_shape = (-1, tf.shape(y_pred)[-1] - 1)
+    yl = tf.boolean_mask(y_pred, y_true > 0)
+    ynl = tf.reshape(
+        tf.boolean_mask(y_pred, y_true <= 0),
+        ynl_shape,
+    )
+    delta = yl - tf.reduce_max(ynl, axis=-1)
+    return delta
+
+
+def _delta_binary(y_true, y_pred):
+    r"""
+    Compute the non-normalized provable robustness factor in binary setup (need to be
+    adjusted with
+    lipschitz constant).
+
+    .. math::
+            \Delta(x) = f(x) \text{ if } l=1, -f(x) \text{ otherwise}
+
+    Args:
+        y_true: true labels must be in {1,0} or in {1,-1} (no label smoothing allowed)
+        y_pred: network predictions
+
+    Returns: non-normalized provable robustness factor
+
+    """
+    y_true = tf.sign(tf.cast(y_true, y_pred.dtype) - 1e-3)
+    return tf.multiply(y_true, y_pred)
+
+
+@register_keras_serializable("deel-lip", "CategoricalProvableRobustAccuracy")
+class CategoricalProvableRobustAccuracy(Loss):
+    def __init__(
+        self,
+        epsilon=36 / 255,
+        lip_const=1.0,
+        disjoint_neurons=True,
+        reduction=Reduction.AUTO,
+        name="CategoricalProvableRobustAccuracy",
+    ):
+        r"""
+
+        The accuracy that can be proved at a given epsilon.
+
+        Args:
+            epsilon: the metric will return the guaranteed accuracy for the radius
+                epsilon
+            lip_const: lipschitz constant of the network
+            disjoint_neurons: must be set to True if your model ends with a
+                FrobeniusDense layer with `disjoint_neurons` set to True. Set to False
+                otherwise
+            reduction: the recution method when training in a multi-gpu / TPU system
+            name: metrics name.
+        """
+        self.lip_const = lip_const
+        self.epsilon = epsilon
+        self.disjoint_neurons = disjoint_neurons
+        if disjoint_neurons:
+            self.certificate_factor = 2 * lip_const
+        else:
+            self.certificate_factor = math.sqrt(2) * lip_const
+        super(CategoricalProvableRobustAccuracy, self).__init__(reduction, name)
+
+    @tf.function
+    def call(self, y_true, y_pred):
+        return tf.cast(
+            (_delta_multiclass(y_true, y_pred) / self.certificate_factor)
+            > self.epsilon,
+            y_pred.dtype,
+        )
+
+    def get_config(self):
+        config = {
+            "epsilon": self.epsilon,
+            "lip_const": self.lip_const,
+            "disjoint_neurons": self.disjoint_neurons,
+        }
+        base_config = super(CategoricalProvableRobustAccuracy, self).get_config()
+        return dict(list(base_config.items()) + list(config.items()))
+
+
+@register_keras_serializable("deel-lip", "BinaryProvableRobustAccuracy")
+class BinaryProvableRobustAccuracy(Loss):
+    def __init__(
+        self,
+        epsilon=36 / 255,
+        lip_const=1.0,
+        reduction=Reduction.AUTO,
+        name="BinaryProvableRobustAccuracy",
+    ):
+        r"""
+
+        The accuracy that can be proved at a given epsilon.
+
+        Args:
+            epsilon: the metric will return the guaranteed accuracy for the radius
+                epsilon
+            lip_const: lipschitz constant of the network
+            reduction: the recution method when training in a multi-gpu / TPU system
+            name: metrics name.
+        """
+        self.lip_const = lip_const
+        self.epsilon = epsilon
+        super(BinaryProvableRobustAccuracy, self).__init__(reduction, name)
+
+    @tf.function
+    def call(self, y_true, y_pred):
+        return tf.cast(
+            (_delta_binary(y_true, y_pred) / self.lip_const) > self.epsilon,
+            y_pred.dtype,
+        )
+
+    def get_config(self):
+        config = {
+            "epsilon": self.epsilon,
+            "lip_const": self.lip_const,
+        }
+        base_config = super(BinaryProvableRobustAccuracy, self).get_config()
+        return dict(list(base_config.items()) + list(config.items()))
+
+
+@register_keras_serializable("deel-lip", "CategoricalProvableAvgRobustness")
+class CategoricalProvableAvgRobustness(Loss):
+    def __init__(
+        self,
+        lip_const=1.0,
+        disjoint_neurons=True,
+        negative_robustness=False,
+        reduction=Reduction.AUTO,
+        name="CategoricalProvableAvgRobustness",
+    ):
+        r"""
+
+        Compute the average provable robustness radius on the dataset.
+
+        .. math::
+            \mathbb{E}_{x \in D}\left[ \frac{\phi\left(\mathcal{M}_f(x)\right)}{
+            L_f}\right]
+
+        :math:`\mathcal{M}_f(x)` is a term that: is positive when x is correctly
+        classified and negative otherwise. In both case the value give the robustness
+        radius around x.
+
+        In the multiclass setup we have:
+
+        .. math::
+            \mathcal{M}_f(x) =f_l(x) - \max_{i \neq l} f_i(x)
+
+        Where :math:`D` is the dataset, :math:`l` is the correct label for x and
+        :math:`L_f` is the lipschitz constant of the network (:math:`L = 2 \times
+        \text{lip_const}` when `disjoint_neurons=True`, :math:`L = \sqrt{2} \times
+        \text{lip_const}` otherwise).
+
+        When `negative_robustness` is set to `True` misclassified elements count as
+        negative robustness (:math:`\phi` act as identity function), when set to
+        `False`,
+        misclassified elements yield a robustness radius of 0 ( :math:`\phi(x)=relu(
+        x)` ). The elements are not ignored when computing the mean in both cases.
+
+        This metric works for labels both in {1,0} and {1,-1}.
+
+        Args:
+            lip_const: lipschitz constant of the network
+            disjoint_neurons: must be set to True is your model ends with a
+                FrobeniusDense layer with `disjoint_neurons` set to True. Set to False
+                otherwise
+            reduction: the recution method when training in a multi-gpu / TPU system
+            name: metrics name.
+        """
+        self.lip_const = lip_const
+        self.disjoint_neurons = disjoint_neurons
+        self.negative_robustness = negative_robustness
+        if disjoint_neurons:
+            self.certificate_factor = 2 * lip_const
+        else:
+            self.certificate_factor = math.sqrt(2) * lip_const
+        if self.negative_robustness:
+            self.delta_correction = lambda delta: delta
+        else:
+            self.delta_correction = tf.nn.relu
+        super(CategoricalProvableAvgRobustness, self).__init__(reduction, name)
+
+    @tf.function
+    def call(self, y_true, y_pred):
+        return (
+            self.delta_correction(_delta_multiclass(y_true, y_pred))
+            / self.certificate_factor
+        )
+
+    def get_config(self):
+        config = {
+            "lip_const": self.lip_const,
+            "disjoint_neurons": self.disjoint_neurons,
+            "negative_robustness": self.negative_robustness,
+        }
+        base_config = super(CategoricalProvableAvgRobustness, self).get_config()
+        return dict(list(base_config.items()) + list(config.items()))
+
+
+@register_keras_serializable("deel-lip", "BinaryProvableAvgRobustness")
+class BinaryProvableAvgRobustness(Loss):
+    def __init__(
+        self,
+        lip_const=1.0,
+        negative_robustness=False,
+        reduction=Reduction.AUTO,
+        name="BinaryProvableAvgRobustness",
+    ):
+        r"""
+
+        Compute the average provable robustness radius on the dataset.
+
+        .. math::
+            \mathbb{E}_{x \in D}\left[ \frac{\phi\left(\mathcal{M}_f(x)\right)}{
+            L_f}\right]
+
+        :math:`\mathcal{M}_f(x)` is a term that: is positive when x is correctly
+        classified and negative otherwise. In both case the value give the robustness
+        radius around x.
+
+        In the binary classification setup we have:
+
+        .. math::
+            \mathcal{M}_f(x) = f(x) \text{ if } l=1, -f(x) \text{otherwise}
+
+        Where :math:`D` is the dataset, :math:`l` is the correct label for x and
+        :math:`L_f` is the lipschitz constant of the network..
+
+        When `negative_robustness` is set to `True` misclassified elements count as
+        negative robustness (:math:`\phi` act as identity function), when set to
+        `False`,
+        misclassified elements yield a robustness radius of 0 ( :math:`\phi(x)=relu(
+        x)` ). The elements are not ignored when computing the mean in both cases.
+
+        This metric works for labels both in {1,0} and {1,-1}.
+
+        Args:
+            lip_const: lipschitz constant of the network
+            reduction: the recution method when training in a multi-gpu / TPU system
+            name: metrics name.
+        """
+        self.lip_const = lip_const
+        self.negative_robustness = negative_robustness
+        if self.negative_robustness:
+            self.delta_correction = lambda delta: delta
+        else:
+            self.delta_correction = tf.nn.relu
+        super(BinaryProvableAvgRobustness, self).__init__(reduction, name)
+
+    @tf.function
+    def call(self, y_true, y_pred):
+        return self.delta_correction(_delta_binary(y_true, y_pred)) / self.lip_const
+
+    def get_config(self):
+        config = {
+            "lip_const": self.lip_const,
+            "negative_robustness": self.negative_robustness,
+        }
+        base_config = super(BinaryProvableAvgRobustness, self).get_config()
+        return dict(list(base_config.items()) + list(config.items()))
diff --git a/doc/source/deel.lip.rst b/doc/source/deel.lip.rst
index c218ce80..85e72fdc 100644
--- a/doc/source/deel.lip.rst
+++ b/doc/source/deel.lip.rst
@@ -13,6 +13,7 @@ Submodules
    deel.lip.initializers
    deel.lip.layers
    deel.lip.losses
+   deel.lip.metrics
    deel.lip.model
    deel.lip.normalizers
    deel.lip.utils


Test patch:
diff --git a/tests/test_metrics.py b/tests/test_metrics.py
new file mode 100644
index 00000000..0e0cda72
--- /dev/null
+++ b/tests/test_metrics.py
@@ -0,0 +1,278 @@
+import unittest
+from collections import defaultdict
+from unittest import TestCase
+import tensorflow as tf
+from tensorflow.keras.models import Sequential, load_model
+from tensorflow.keras.layers import Input, Dense
+from tensorflow.keras.optimizers import SGD
+from deel.lip.metrics import (
+    CategoricalProvableRobustAccuracy,
+    BinaryProvableRobustAccuracy,
+    CategoricalProvableAvgRobustness,
+    BinaryProvableAvgRobustness,
+)
+import numpy as np
+import os
+
+
+def check_serialization(nb_class, loss, nb_classes):
+    n = 255
+    x = tf.random.uniform((n, 42))
+    y = tf.one_hot(
+        tf.convert_to_tensor(np.random.randint(nb_classes, size=n)), nb_classes
+    )
+    m = Sequential([Input(42), Dense(nb_class)])
+    m.compile(optimizer=SGD(), loss=loss)
+    l1 = m.evaluate(x, y)
+    name = loss.__class__.__name__
+    path = os.path.join("logs", "losses", name)
+    m.save(path)
+    m2 = load_model(path, compile=True)
+    l2 = m2.evaluate(x, y)
+    return l1, l2
+
+
+class Test(TestCase):
+    def test_serialization(self):
+        pra = CategoricalProvableRobustAccuracy(1, 1.0, disjoint_neurons=False)
+        l1, l2 = check_serialization(10, pra, 10)
+        self.assertEqual(
+            l1,
+            l2,
+            "serialization changed loss value for CategoricalProvableRobustAccuracy",
+        )
+        pra = BinaryProvableRobustAccuracy(1, 1.0)
+        l1, l2 = check_serialization(10, pra, 10)
+        self.assertEqual(
+            l1, l2, "serialization changed loss value for BinaryProvableRobustAccuracy"
+        )
+        par = CategoricalProvableAvgRobustness(
+            2.0, disjoint_neurons=False, negative_robustness=True
+        )
+        l1, l2 = check_serialization(10, par, 10)
+        self.assertEqual(
+            l1,
+            l2,
+            "serialization changed loss value for CategoricalProvableAvgRobustness",
+        )
+        par = BinaryProvableAvgRobustness(2.0, negative_robustness=True)
+        l1, l2 = check_serialization(10, par, 10)
+        self.assertEqual(
+            l1, l2, "serialization changed loss value for BinaryProvableAvgRobustness"
+        )
+
+    def test_provable_vs_adjusted(self):
+        n = 500
+        x = tf.random.normal((n, 10), mean=0.0, stddev=0.1)
+        y = tf.one_hot(tf.convert_to_tensor(np.random.randint(10, size=n)), 10)
+        for lip_cst in [1, 0.5, 2.0]:
+            for disjoint in [True, False]:
+                pr = CategoricalProvableAvgRobustness(
+                    lip_cst, disjoint, negative_robustness=False
+                )
+                ar = CategoricalProvableAvgRobustness(
+                    lip_cst, disjoint, negative_robustness=True
+                )
+                l1 = pr(y, y).numpy()
+                l2 = ar(y, y).numpy()
+                self.assertAlmostEqual(
+                    l1,
+                    l2,
+                    4,
+                    msg="provable and adjusted "
+                    "robustness must give same "
+                    "values when y_true==y_pred",
+                )
+                l1 = pr(y, x).numpy()
+                l2 = ar(y, x).numpy()
+                self.assertNotAlmostEqual(
+                    l1,
+                    l2,
+                    4,
+                    msg="provable and adjusted robustness must give different values"
+                    " when y_true!=y_pred",
+                )
+
+        x = tf.random.normal((n, 10), mean=0.0, stddev=0.1)
+        y = tf.one_hot(tf.convert_to_tensor(np.random.randint(10, size=n)), 10)
+        for lip_cst in [1.0, 0.5, 2.0]:
+            bpr = BinaryProvableAvgRobustness(lip_cst, negative_robustness=False)
+            bar = BinaryProvableAvgRobustness(lip_cst, negative_robustness=True)
+            l1 = bpr(y, y).numpy()
+            l2 = bar(y, y).numpy()
+            self.assertAlmostEqual(
+                l1,
+                l2,
+                4,
+                msg="provable and adjusted robustness must give same values when y_"
+                "true==y_pred",
+            )
+            l1 = pr(y, x).numpy()
+            l2 = ar(y, x).numpy()
+            self.assertNotAlmostEqual(
+                l1,
+                l2,
+                4,
+                msg="provable and adjusted robustness must give different values when "
+                "y_true!=y_pred",
+            )
+
+    def test_data_format(self):
+        pr = CategoricalProvableAvgRobustness(1.0, True, negative_robustness=False)
+        ar = CategoricalProvableAvgRobustness(1.0, True, negative_robustness=True)
+        n = 500
+        # check in multiclass
+        for metric in [pr, ar]:
+            y_pred = tf.random.normal((n, 10), mean=0.0, stddev=0.1)
+            y_true = tf.one_hot(tf.convert_to_tensor(np.random.randint(10, size=n)), 10)
+            metrics_values = []
+            for neg_val in [0.0, -1.0]:
+                y_pred_case = y_pred
+                y_true_case = tf.where(y_true > 0, 1.0, neg_val)
+                metrics_values.append(metric(y_true_case, y_pred_case).numpy())
+            self.assertTrue(
+                all([m == metrics_values[0] for m in metrics_values]),
+                "changing the data format must not change the metric value",
+            )
+        # check in binary
+        bpr = BinaryProvableAvgRobustness(1.0, negative_robustness=False)
+        bar = BinaryProvableAvgRobustness(1.0, negative_robustness=True)
+        for metric in [bpr, bar]:
+            y_pred = tf.random.normal((n,), mean=0.0, stddev=0.1)
+            y_true = tf.random.normal((n,), mean=0.0, stddev=0.1)
+            metrics_values = []
+            for neg_val in [0.0, -1.0]:
+                y_pred_case = y_pred
+                y_true_case = tf.where(y_true > 0, 1.0, neg_val)
+                metrics_values.append(metric(y_true_case, y_pred_case).numpy())
+            self.assertTrue(
+                all([m == metrics_values[0] for m in metrics_values]),
+                "changing the data format must not change the metric value",
+            )
+
+    def test_disjoint_neurons(self):
+        pr = CategoricalProvableAvgRobustness(1.0, False, negative_robustness=False)
+        ar = CategoricalProvableAvgRobustness(1.0, False, negative_robustness=True)
+        pdr = CategoricalProvableAvgRobustness(1.0, True, negative_robustness=False)
+        adr = CategoricalProvableAvgRobustness(1.0, True, negative_robustness=True)
+        n = 500
+        # check in multiclass
+        metrics_values = defaultdict(list)
+        y_pred = tf.random.normal((n, 10), mean=0.0, stddev=1.0)
+        y_true = tf.one_hot(tf.convert_to_tensor(np.random.randint(10, size=n)), 10)
+        for neg_val in [0.0, -1.0]:
+            y_pred_case = y_pred
+            y_true_case = tf.where(y_true > 0, 1.0, neg_val)
+            metrics_values["standard"].append(pdr(y_true_case, y_pred_case).numpy())
+            metrics_values["standard"].append(
+                pr(y_true_case, y_pred_case * (np.sqrt(2) / 2.0)).numpy()
+            )
+            metrics_values["negative_rob"].append(adr(y_true_case, y_pred_case).numpy())
+            metrics_values["negative_rob"].append(
+                ar(y_true_case, y_pred_case * (np.sqrt(2) / 2.0)).numpy()
+            )
+        print(metrics_values)
+        for metric_type in ["standard", "negative_rob"]:
+            for m in metrics_values[metric_type]:
+                self.assertAlmostEqual(
+                    m,
+                    metrics_values[metric_type][0],
+                    4,
+                    "the loss does not account disjoint_neuron properly in multiclass",
+                )
+        # check in binary
+        pr = BinaryProvableAvgRobustness(1.0, negative_robustness=False)
+        ar = BinaryProvableAvgRobustness(1.0, negative_robustness=True)
+        metrics_values = defaultdict(list)
+        y_pred = tf.random.normal((n,), mean=0.0, stddev=0.1)
+        y_true = tf.random.normal((n,), mean=0.0, stddev=0.1)
+        for neg_val in [0.0, -1.0]:
+            y_pred_case = y_pred
+            y_true_case = tf.where(y_true > 0, 1.0, neg_val)
+            # no sqrt(2)/2 here as the corrective factor works only in multiclass
+            metrics_values["standard"].append(pr(y_true_case, y_pred_case).numpy())
+            metrics_values["negative_rob"].append(ar(y_true_case, y_pred_case).numpy())
+        self.assertTrue(
+            all(
+                [m == metrics_values["standard"][0] for m in metrics_values["standard"]]
+            ),
+            "the loss does not account disjoint_neuron properly in binary",
+        )
+        self.assertTrue(
+            all(
+                [
+                    m == metrics_values["negative_rob"][0]
+                    for m in metrics_values["negative_rob"]
+                ]
+            ),
+            "the loss does not account disjoint_neuron properly in binary",
+        )
+
+    def test_hardcoded_values(self):
+        pr = CategoricalProvableRobustAccuracy(0.25, 1.0, False)
+        ar = CategoricalProvableAvgRobustness(1.0, False)
+        y_pred = tf.convert_to_tensor(
+            [
+                [1.0, 0.0, 0.0],  # good class & over 0.25
+                [0.1, 0.0, 0.0],  # good class & below 0.25
+                [0.0, 0.0, 1.0],  # wrong class & over 0.25
+                [0.0, 0.0, 0.1],  # wrong class & below 0.25
+            ]
+        )
+        y_true = tf.convert_to_tensor(
+            [
+                [1.0, 0.0, 0.0],
+                [1.0, 0.0, 0.0],
+                [0.0, 1.0, 0.0],
+                [0.0, 1.0, 0.0],
+            ]
+        )
+        self.assertEqual(
+            pr(y_true, y_pred).numpy(),
+            0.25,
+            "ProvableRobustAccuracy did not "
+            "get the expected result with "
+            "hardcoded values ",
+        )
+        self.assertAlmostEqual(
+            ar(y_true, y_pred).numpy(),
+            0.25 * 1.1 / np.sqrt(2),
+            5,
+            "ProvableAvgRobustness did not "
+            "get the expected result with "
+            "hardcoded values ",
+        )
+        bpr = BinaryProvableRobustAccuracy(0.25, 1.0)
+        bar = BinaryProvableAvgRobustness(1.0, negative_robustness=False)
+        y_pred = tf.convert_to_tensor(
+            [
+                1.0,  # good class & over 0.25
+                0.1,  # good class & below 0.25
+                1.0,  # wrong class & over 0.25
+                0.1,  # wrong class & below 0.25
+                -1.0,
+                -0.1,
+                -1.0,
+                -0.1,
+            ]
+        )
+        y_true = tf.convert_to_tensor([1.0, 1.0, -1.0, -1.0, -1.0, -1.0, 1.0, 1.0])
+        self.assertEqual(
+            bpr(y_true, y_pred).numpy(),
+            0.25,
+            "BinaryProvableRobustAccuracy did not "
+            "get the expected result with "
+            "hardcoded values ",
+        )
+        self.assertAlmostEqual(
+            bar(y_true, y_pred).numpy(),
+            0.125 * (1.1 * 2),
+            5,
+            "BinaryProvableAvgRobustness did not "
+            "get the expected result with "
+            "hardcoded values ",
+        )
+
+
+if __name__ == "__main__":
+    unittest.main()