Feat/keras3 test uniformization

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

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):

setup.cfg

@@ -14,12 +14,13 @@ envlist =
 
 [testenv]
 deps =
+    pytest
     tf216: tensorflow ~= 2.16.0
     tf217: tensorflow ~= 2.17.0
 
 
 commands =
-    python -m unittest
+    pytest tests
 
 [testenv:py{39,310,311}-lint]
 skip_install = true