Merge pull request #96 from deel-ai/feat/keras3_test_uniformization

Feat/keras3 test uniformization

Author: franckma31

CONTRIBUTING.md

@@ -32,7 +32,7 @@ $ make test
 
 This command will:
 - check your code with black PEP-8 formatter and flake8 linter.
-- run `unittest` on the `tests/` folder with different Python and TensorFlow versions.
+- run `pytest` on the `tests/` folder with different Python and TensorFlow versions.
 
 
 ## Submitting your changes

README.md

@@ -1,8 +1,8 @@
 <div align="center">
         <picture>
-                <source media="(prefers-color-scheme: dark)" srcset="./docs/assets/logo_white.svg">
-                <source media="(prefers-color-scheme: light)" srcset="./docs/assets/logo.svg">
-                <img alt="Library Banner" src="./docs/assets/logo.svg">
+                <source media="(prefers-color-scheme: dark)" srcset="./docs/assets/banner_dark_deellip.png">
+                <source media="(prefers-color-scheme: light)" srcset="./docs/assets/banner_light_deellip.png">
+                <img alt="DEEL-LIP Banner" src="./docs/assets/banner_light_deellip.png">
         </picture>
 </div>
 <br>
@@ -12,10 +12,10 @@
         <img src="https://img.shields.io/pypi/pyversions/deel-lip.svg">
     </a>
     <a href="https://github.com/deel-ai/deel-lip/actions/workflows/python-linters.yml">
-        <img alt="PyLint" src="https://github.com/deel-ai/deel-lip/actions/workflows/python-linters.yml/badge.svg?branch=master">
+        <img alt="PyLint" src="https://github.com/deel-ai/deel-lip/actions/workflows/python-linters.yml/badge.svg?branch=keras3">
     </a>
     <a href="https://github.com/deel-ai/deel-lip/actions/workflows/python-tests.yml">
-        <img alt="Tox" src="https://github.com/deel-ai/deel-lip/actions/workflows/python-linters.yml/badge.svg?branch=master">
+        <img alt="Tox" src="https://github.com/deel-ai/deel-lip/actions/workflows/python-linters.yml/badge.svg?branch=keras3">
     </a>
     <a href="https://pypi.org/project/deel-lip">
         <img alt="Pypi" src="https://img.shields.io/pypi/v/deel-lip.svg">
@@ -38,19 +38,15 @@ has many applications ranging from adversarial robustness to Wasserstein
 distance estimation.
 
 This library provides an efficient implementation of **k-Lispchitz
-layers for keras**.
+layers for Keras 3**.
 
 > [!CAUTION]
-> **Incompatibility with TensorFlow >= 2.16 and Keras 3**
+> **This branch is a major update designed for compatibility with TensorFlow ≥ 2.16 and Keras 3**
 >
-> Due to significant changes introduced in TensorFlow version 2.16 and Keras 3, this
-> package is currently incompatible with TensorFlow versions 2.16 and above. Users are
-> advised to use TensorFlow versions lower than 2.16 to ensure compatibility and proper
+> Due to significant changes introduced in TensorFlow 2.16 and Keras 3, backward compatibility 
+> with older TensorFlow versions can no longer be maintained. Users requiring an older version 
+> of TensorFlow are encouraged to clone the master branch to ensure compatibility and proper
 > functionality of this package.
->
-> We are actively working on updating the package to support Keras 3. Please stay tuned
-> for updates. For now, make sure to install an earlier version of TensorFlow by
-> specifying it in your environment.
 
 ## 📚 Table of contents
 
@@ -69,7 +65,7 @@ layers for keras**.
 You can install ``deel-lip`` directly from pypi:
 
 ```python
-pip install deel-lip
+pip install git+https://github.com/deel-ai/deel-lip.git@keras3
 ```
 
 In order to use ``deel-lip``, you also need a [valid tensorflow
@@ -80,12 +76,12 @@ supports tensorflow versions 2.x.
 
 | **Tutorial Name**           | Notebook                                                                                                                                                           |
 | :-------------------------- | :----------------------------------------------------------------------------------------------------------------------------------------------------------------: |
-| Getting Started 1 - Creating a 1-Lipschitz neural network | [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/deel-ai/deel-lip/blob/master/docs/notebooks/Getting_started_1.ipynb)            |
-| Getting Started 2 - Training an adversarially robust 1-Lipschitz neural network | [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/deel-ai/deel-lip/blob/master/docs/notebooks/Getting_started_2.ipynb)            |
-| Wasserstein distance estimation on toy example | [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/deel-ai/deel-lip/blob/master/docs/notebooks/demo1.ipynb) |
-| HKR Classifier on toy dataset | [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/deel-ai/deel-lip/blob/master/docs/notebooks/demo2.ipynb) |
-| HKR classifier on MNIST dataset | [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/deel-ai/deel-lip/blob/master/docs/notebooks/demo3.ipynb) |
-| HKR multiclass and fooling | [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/deel-ai/deel-lip/blob/master/docs/notebooks/demo4.ipynb) |
+| Getting Started 1 - Creating a 1-Lipschitz neural network | [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/deel-ai/deel-lip/blob/keras3/docs/notebooks/Getting_started_1.ipynb)            |
+| Getting Started 2 - Training an adversarially robust 1-Lipschitz neural network | [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/deel-ai/deel-lip/blob/keras3/docs/notebooks/Getting_started_2.ipynb)            |
+| Wasserstein distance estimation on toy example | [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/deel-ai/deel-lip/blob/keras3/docs/notebooks/demo1.ipynb) |
+| HKR Classifier on toy dataset | [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/deel-ai/deel-lip/blob/keras3/docs/notebooks/demo2.ipynb) |
+| HKR classifier on MNIST dataset | [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/deel-ai/deel-lip/blob/keras3/docs/notebooks/demo3.ipynb) |
+| HKR multiclass and fooling | [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/deel-ai/deel-lip/blob/keras3/docs/notebooks/demo4.ipynb) |
 
 
 ## 📦 What's Included
@@ -127,7 +123,8 @@ More from the DEEL project:
 
 - [Xplique](https://github.com/deel-ai/xplique) a Python library exclusively dedicated to explaining neural networks.
 - [Influenciae](https://github.com/deel-ai/influenciae) Python toolkit dedicated to computing influence values for the discovery of potentially problematic samples in a dataset.
-- [deel-torchlip](https://github.com/deel-ai/deel-torchlip) a Python library for training k-Lipschitz neural networks on PyTorch.
+- [deel-TorchLip](https://github.com/deel-ai/deel-torchlip) a Python library for training k-Lipschitz neural networks on PyTorch.
+- [Oodeel](https://github.com/deel-ai/oodeel) a Python library for post-hoc deep OOD (Out-of-Distribution) detection on already trained neural network image classifiers
 - [DEEL White paper](https://arxiv.org/abs/2103.10529) a summary of the DEEL team on the challenges of certifiable AI and the role of data quality, representativity and explainability for this purpose.
 
 ## 🙏 Acknowledgments

deel/lip/layers/convolutional.py

@@ -441,6 +441,10 @@ def __init__(
             raise ValueError("SpectralConv2DTranspose does not support dilation rate")
         if self.padding != "same":
             raise ValueError("SpectralConv2DTranspose only supports padding='same'")
+        if self.output_padding is not None:
+            raise ValueError(
+                "SpectralConv2DTranspose only supports output_padding=None"
+            )
         self.set_klip_factor(k_coef_lip)
         self.u = None
         self.sig = None

deel/lip/layers/pooling.py

@@ -22,6 +22,7 @@
 """
 
 import numpy as np
+from typing import Tuple
 import keras
 import keras.ops as K
 from keras.saving import register_keras_serializable
@@ -91,6 +92,7 @@ def __init__(
             data_format=data_format,
             **kwargs,
         )
+        self.built = False
         self.set_klip_factor(k_coef_lip)
         self._kwargs = kwargs
 
@@ -181,6 +183,7 @@ def __init__(
             data_format=data_format,
             **kwargs,
         )
+        self.built = False
         self.set_klip_factor(k_coef_lip)
         self.eps_grad_sqrt = eps_grad_sqrt
         self._kwargs = kwargs
@@ -246,6 +249,7 @@ def __init__(self, data_format=None, k_coef_lip=1.0, eps_grad_sqrt=1e-6, **kwarg
         super(ScaledGlobalL2NormPooling2D, self).__init__(
             data_format=data_format, **kwargs
         )
+        self.built = False
         self.set_klip_factor(k_coef_lip)
         self.eps_grad_sqrt = eps_grad_sqrt
         self._kwargs = kwargs
@@ -308,6 +312,7 @@ def __init__(self, data_format=None, k_coef_lip=1.0, **kwargs):
         super(ScaledGlobalAveragePooling2D, self).__init__(
             data_format=data_format, **kwargs
         )
+        self.built = False
         self.set_klip_factor(k_coef_lip)
         self._kwargs = kwargs
 
@@ -363,32 +368,44 @@ def __init__(
             **kwargs: params passed to the Layers constructor
         """
         super(InvertibleDownSampling, self).__init__(name=name, dtype=dtype, **kwargs)
-        self.pool_size = pool_size
         self.data_format = data_format
 
-    def call(self, inputs):
-        if self.data_format == "channels_last":
-            return K.concatenate(
-                [
-                    inputs[
-                        :, i :: self.pool_size[0], j :: self.pool_size[1], :
-                    ]  # for now we handle only channels last
-                    for i in range(self.pool_size[0])
-                    for j in range(self.pool_size[1])
-                ],
-                axis=-1,
-            )
+        ndims = 2
+        ks: Tuple[int, ...]
+        if isinstance(pool_size, int):
+            ks = (pool_size,) * ndims
         else:
-            return K.concatenate(
-                [
-                    inputs[
-                        :, :, i :: self.pool_size[0], j :: self.pool_size[1]
-                    ]  # for now we handle only channels last
-                    for i in range(self.pool_size[0])
-                    for j in range(self.pool_size[1])
-                ],
-                axis=1,
+            ks = tuple(pool_size)
+
+        if len(ks) != ndims:
+            raise ValueError(
+                f"Expected {ndims}-dimensional pool_size, but "
+                f"got {len(ks)}-dimensional instead"
+            )
+        self.pool_size = ks
+
+    def call(self, inputs):
+        if self.data_format == "channels_first":
+            # convert to channels_first
+            inputs = K.transpose(inputs, [0, 2, 3, 1])
+        # from shape (bs, w*pw, h*ph, c) to (bs, w, h, c, pw, ph)
+        input_shape = K.shape(inputs)
+        w, h, c_in = input_shape[1], input_shape[2], input_shape[3]
+        pw, ph = self.pool_size
+        wo = w // pw
+        ho = h // ph
+        inputs = K.reshape(inputs, (-1, wo, pw, h, c_in))
+        inputs = K.reshape(inputs, (-1, wo, pw, ho, ph, c_in))
+        inputs = K.transpose(
+            inputs, [0, 1, 3, 5, 2, 4]
+        )  # (bs, wo, pw, ho, ph, c) -> (bs, wo, ho, c, pw, ph)
+        inputs = K.reshape(inputs, (-1, wo, ho, c_in * pw * ph))
+
+        if self.data_format == "channels_first":
+            inputs = K.transpose(
+                inputs, [0, 3, 1, 2]  # (bs, w, h, c*pw*ph) -> (bs, c*pw*ph, w, h) ->
             )
+        return inputs
 
     def get_config(self):
         config = {
@@ -427,9 +444,22 @@ def __init__(
             **kwargs: params passed to the Layers constructor
         """
         super(InvertibleUpSampling, self).__init__(name=name, dtype=dtype, **kwargs)
-        self.pool_size = pool_size
         self.data_format = data_format
 
+        ndims = 2
+        ks: Tuple[int, ...]
+        if isinstance(pool_size, int):
+            ks = (pool_size,) * ndims
+        else:
+            ks = tuple(pool_size)
+
+        if len(ks) != ndims:
+            raise ValueError(
+                f"Expected {ndims}-dimensional pool_size, but "
+                f"got {len(ks)}-dimensional instead"
+            )
+        self.pool_size = ks
+
     def call(self, inputs):
         if self.data_format == "channels_first":
             # convert to channels_first
@@ -439,12 +469,12 @@ def call(self, inputs):
         w, h, c_in = input_shape[1], input_shape[2], input_shape[3]
         pw, ph = self.pool_size
         c = c_in // (pw * ph)
-        inputs = K.reshape(inputs, (-1, w, h, pw, ph, c))
+        inputs = K.reshape(inputs, (-1, w, h, c, pw, ph))
         inputs = K.transpose(
             K.reshape(
                 K.transpose(
-                    inputs, [0, 5, 2, 4, 1, 3]
-                ),  # (bs, w, h, pw, ph, c) -> (bs, c, w, pw, h, ph)
+                    inputs, [0, 3, 2, 5, 1, 4]
+                ),  # (bs, w, h, c, pw, ph) -> (bs, c, w, pw, h, ph)
                 (-1, c, w, pw, h * ph),
             ),  # (bs, c, w, pw, h, ph) -> (bs, c, w, pw, h*ph) merge last axes
             [

deel/lip/utils.py

@@ -61,35 +61,36 @@ def evaluate_lip_const_gen(
 
 def evaluate_lip_const(model: keras.Model, x, eps=1e-4, seed=None):
     """
-    Evaluate the Lipschitz constant of a model, with the naive method.
-    Please note that the estimation of the lipschitz constant is done locally around
-    input sample. This may not correctly estimate the behaviour in the whole domain.
+    Evaluate the Lipschitz constant of a model using the Jacobian of the model.
+    The estimation is done locally around input samples.
 
     Args:
-        model: built keras model used to make predictions
-        x: inputs used to compute the lipschitz constant
-        eps (float): magnitude of noise to add to input in order to compute the constant
-        seed (int): seed used when generating the noise ( can be set to None )
+        model (Model): A built Keras model used to make predictions.
+        x (np.ndarray): Input samples used to compute the Lipschitz constant.
 
     Returns:
-        float: the empirically evaluated lipschitz constant. The computation might also
-            be inaccurate in high dimensional space.
-
+        float: The empirically evaluated Lipschitz constant.
     """
-    y_pred = model.predict(x)
-    # x = np.repeat(x, 100, 0)
-    # y_pred = np.repeat(y_pred, 100, 0)
-    x_var = x + keras.random.uniform(
-        shape=x.shape, minval=eps * 0.25, maxval=eps, seed=seed
-    )
-    y_pred_var = model.predict(x_var)
-    dx = x - x_var
-    dfx = y_pred - y_pred_var
-    ndx = K.sqrt(K.sum(K.square(dx), axis=range(1, len(x.shape))))
-    ndfx = K.sqrt(K.sum(K.square(dfx), axis=range(1, len(y_pred.shape))))
-    lip_cst = K.max(ndfx / ndx)
-    print(f"lip cst: {lip_cst:.3f}")
-    return lip_cst
+    batch_size = x.shape[0]
+    x = keras.ops.convert_to_tensor(x, dtype=model.inputs[0].dtype)
+
+    if keras.config.backend() == "tensorflow":
+        import tensorflow as tf
+
+        with tf.GradientTape() as tape:
+            tape.watch(x)
+            y_pred = model(x, training=False)
+        batch_jacobian = tape.batch_jacobian(y_pred, x)
+    else:
+        assert False, "Only tensorflow backend is supported for now."
+    # Flatten input/output dimensions for spectral norm computation
+    xdim = keras.ops.prod(keras.ops.shape(x)[1:])
+    ydim = keras.ops.prod(keras.ops.shape(y_pred)[1:])
+    batch_jacobian = keras.ops.reshape(batch_jacobian, (batch_size, ydim, xdim))
+
+    # Compute spectral norm of the Jacobians and return the maximum
+    spectral_norms = keras.ops.linalg.norm(batch_jacobian, ord=2, axis=[-2, -1])
+    return keras.ops.max(spectral_norms).numpy()
 
 
 def _padding_circular(x, circular_paddings):