Merge pull request #91 from deel-ai/fix/upgrade_tf_2.17

[BREAKING CHANGE] deel-lip upgrade to Keras 3.0

TensorFlow 2.16 version and above come with Keras 3: a lot of major changes and incompatibilities have been introduced with previous TF versions.

We then propose a new deel-lip version 2.0.0 compatible with TensorFlow 2.16+ and Keras 3.
This PR has two sets of commits:

- a set of "fix" commits that resolves bugs introduced by TF2.16+ and Keras 3, compared to TF 2.15 and Keras 2. It is mainly API changes of tf.keras functions and classes.
- a set of "feat" commits that removes all direct dependencies to TensorFlow and only use Keras 3 API. The new deel-lip code is now only based on Keras (with TensorFlow backend) but there is no direct calls to TensorFlow API anymore.

Author: cofri

.github/workflows/python-linters.yml

@@ -13,7 +13,7 @@ jobs:
     strategy:
       max-parallel: 4
       matrix:
-        python-version: [3.7, "3.10"]
+        python-version: ["3.10"]
 
     steps:
     - uses: actions/checkout@v4

.github/workflows/python-tests.yml

@@ -16,14 +16,8 @@ jobs:
       max-parallel: 4
       matrix:
         include:
-          - python-version: 3.7
-            tf-version: 2.3
-          - python-version: 3.9
-            tf-version: 2.7
           - python-version: "3.10"
-            tf-version: 2.11
-          - python-version: "3.10"
-            tf-version: 2.15
+            tf-version: 2.17
 
     steps:
     - uses: actions/checkout@v4

deel/lip/VERSION

@@ -1 +1 @@
-1.5.0
+2.0.0

deel/lip/__init__.py

@@ -25,7 +25,7 @@
 --------
 
 DEEL-LIP provides a simple interface to build and train Lipschitz-constrained neural
-networks based on TensorFlow/Keras framework.
+networks based on Keras framework.
 """
 from os import path
 

deel/lip/callbacks.py

@@ -6,11 +6,12 @@
 This module contains callbacks that can be added to keras training process.
 """
 import os
-from typing import Optional, Dict, Iterable
+from typing import Dict, Iterable, Optional
 
-import tensorflow as tf
-from tensorflow.keras.callbacks import Callback
+import keras.ops as K
 import numpy as np
+from keras.callbacks import Callback
+
 from .layers import Condensable
 
 
@@ -91,6 +92,8 @@ def __init__(
         assert what in {"max", "all"}
         self.what = what
         self.logdir = logdir
+        import tensorflow as tf
+
         self.file_writer = tf.summary.create_file_writer(
             os.path.join(logdir, "metrics")
         )
@@ -101,6 +104,8 @@ def __init__(
         super().__init__()
 
     def _monitor(self, step):
+        import tensorflow as tf
+
         step = self.params["steps"] * self.epochs + step
         for layer_name in self.monitored_layers:
             layer = self.model.get_layer(layer_name)
@@ -113,11 +118,12 @@ def _monitor(self, step):
             elif hasattr(layer, self.target):
                 kernel = getattr(layer, self.target)
                 w_shape = kernel.shape.as_list()
-                sigmas = tf.linalg.svd(
-                    tf.keras.backend.reshape(kernel, [-1, w_shape[-1]]),
+                # TODO: compute_uv=False in next Keras version (3.6.0)
+                sigmas = K.svd(
+                    K.reshape(kernel, [-1, w_shape[-1]]),
                     full_matrices=False,
-                    compute_uv=False,
-                ).numpy()
+                    compute_uv=True,
+                )[1].numpy()
                 sig = sigmas[0]
             else:
                 raise RuntimeWarning(
@@ -176,7 +182,7 @@ def __init__(self, param_name, fp, xp, step=0):
 
         Args:
             param_name (str): name of the parameter of the loss to tune. Must be a
-                tf.Variable.
+                keras.Variable.
             fp (list): values of the loss parameter as steps given by the xp.
             xp (list): step where the parameter equals fp.
             step (int): step value, for serialization/deserialization purposes.
@@ -215,7 +221,7 @@ def __init__(self, param_name, rate=1):
 
     def on_epoch_end(self, epoch: int, logs=None):
         if epoch % self.rate == 0:
-            tf.print(
+            print(
                 "\n",
                 self.model.loss.name,
                 self.param_name,

deel/lip/compute_layer_sv.py

@@ -14,8 +14,8 @@
 It returns a dictionary indicating for each layer name a tuple (min sv, max sv).
 """
 
+import keras
 import numpy as np
-import tensorflow as tf
 
 from .layers import Condensable, GroupSort, MaxMin
 from .layers.unconstrained import PadConv2D
@@ -27,7 +27,7 @@ def _compute_sv_dense(layer, input_sizes=None):
     The singular values are computed using the SVD decomposition of the weight matrix.
 
     Args:
-        layer (tf.keras.Layer): the Dense layer.
+        layer (keras.Layer): the Dense layer.
         input_sizes (tuple, optional): unused here.
 
     Returns:
@@ -46,16 +46,14 @@ def _generate_conv_matrix(layer, input_sizes):
     dirac input.
 
     Args:
-        layer (tf.keras.Layer): the convolutional layer to convert to dense.
+        layer (keras.Layer): the convolutional layer to convert to dense.
         input_sizes (tuple): the input shape of the layer (with batch dimension as first
             element).
 
     Returns:
         np.array: the equivalent matrix of the convolutional layer.
     """
-    single_layer_model = tf.keras.models.Sequential(
-        [tf.keras.layers.Input(input_sizes[1:]), layer]
-    )
+    single_layer_model = keras.Sequential([keras.Input(input_sizes[1:]), layer])
     dirac_inp = np.zeros((input_sizes[2],) + input_sizes[1:])  # Line by line generation
     in_size = input_sizes[1] * input_sizes[2]
     channel_in = input_sizes[-1]
@@ -69,8 +67,8 @@ def _generate_conv_matrix(layer, input_sizes):
                 w_eqmatrix = np.zeros(
                     (in_size * channel_in, np.prod(out_pred.shape[1:]))
                 )
-            w_eqmatrix[start_index : (start_index + input_sizes[2]), :] = tf.reshape(
-                out_pred, (input_sizes[2], -1)
+            w_eqmatrix[start_index : (start_index + input_sizes[2]), :] = (
+                keras.ops.reshape(out_pred, (input_sizes[2], -1))
             )
             dirac_inp = 0.0 * dirac_inp
             start_index += input_sizes[2]
@@ -86,7 +84,7 @@ def _compute_sv_conv2d_layer(layer, input_sizes):
     the weight matrix.
 
     Args:
-        layer (tf.keras.Layer): the convolutional layer.
+        layer (keras.Layer): the convolutional layer.
         input_sizes (tuple): the input shape of the layer (with batch dimension as first
             element).
 
@@ -103,14 +101,14 @@ def _compute_sv_activation(layer, input_sizes=None):
 
     Warning: This is not singular values for non-linear functions but gradient norm.
     """
-    if isinstance(layer, tf.keras.layers.Activation):
-        function2SV = {tf.keras.activations.relu: (0, 1)}
+    if isinstance(layer, keras.layers.Activation):
+        function2SV = {keras.activations.relu: (0, 1)}
         if layer.activation in function2SV.keys():
             return function2SV[layer.activation]
         else:
             return (None, None)
     layer2SV = {
-        tf.keras.layers.ReLU: (0, 1),
+        keras.layers.ReLU: (0, 1),
         GroupSort: (1, 1),
         MaxMin: (1, 1),
     }
@@ -145,25 +143,25 @@ def compute_layer_sv(layer, supplementary_type2sv={}):
     ReLU, Activation, and deel-lip layers)
 
     Args:
-        layer (tf.keras.layers.Layer): a single tf.keras.layer
+        layer (keras.layers.Layer): a single keras.layer
         supplementary_type2sv (dict, optional): a dictionary linking new layer type with
             user-defined function to compute the singular values. Defaults to {}.
     Returns:
         tuple: a 2-tuple with lowest and largest singular values.
     """
     default_type2sv = {
-        tf.keras.layers.Conv2D: _compute_sv_conv2d_layer,
-        tf.keras.layers.Conv2DTranspose: _compute_sv_conv2d_layer,
+        keras.layers.Conv2D: _compute_sv_conv2d_layer,
+        keras.layers.Conv2DTranspose: _compute_sv_conv2d_layer,
         PadConv2D: _compute_sv_conv2d_layer,
-        tf.keras.layers.Dense: _compute_sv_dense,
-        tf.keras.layers.ReLU: _compute_sv_activation,
-        tf.keras.layers.Activation: _compute_sv_activation,
+        keras.layers.Dense: _compute_sv_dense,
+        keras.layers.ReLU: _compute_sv_activation,
+        keras.layers.Activation: _compute_sv_activation,
         GroupSort: _compute_sv_activation,
         MaxMin: _compute_sv_activation,
-        tf.keras.layers.Add: _compute_sv_add,
-        tf.keras.layers.BatchNormalization: _compute_sv_bn,
+        keras.layers.Add: _compute_sv_add,
+        keras.layers.BatchNormalization: _compute_sv_bn,
     }
-    input_shape = layer.input_shape
+    input_shape = layer.input.shape if hasattr(layer.input, "shape") else None
     if isinstance(layer, Condensable):
         layer.condense()
         layer = layer.vanilla_export()
@@ -179,7 +177,7 @@ def compute_model_sv(model, supplementary_type2sv={}):
     """Compute the largest and lowest singular values of all layers in a model.
 
     Args:
-        model (tf.keras.Model): a tf.keras Model or Sequential.
+        model (keras.Model): a keras Model or Sequential.
         supplementary_type2sv (dict, optional): a dictionary linking new layer type
             with user defined function to compute the min and max singular values.
 
@@ -188,7 +186,7 @@ def compute_model_sv(model, supplementary_type2sv={}):
     """
     list_sv = []
     for layer in model.layers:
-        if isinstance(layer, tf.keras.Model):
+        if isinstance(layer, keras.Model):
             list_sv.append((layer.name, (None, None)))
             list_sv += compute_model_sv(layer, supplementary_type2sv)
         else:

deel/lip/constraints.py

@@ -6,16 +6,15 @@
 This module contains extra constraint objects. These object can be added as params to
 regular layers.
 """
-import tensorflow as tf
-from tensorflow.keras import backend as K
-from tensorflow.keras.constraints import Constraint
+import keras.ops as K
+from keras.constraints import Constraint
 from .normalizers import (
     reshaped_kernel_orthogonalization,
     DEFAULT_EPS_SPECTRAL,
     DEFAULT_EPS_BJORCK,
     DEFAULT_BETA_BJORCK,
 )
-from tensorflow.keras.utils import register_keras_serializable
+from keras.saving import register_keras_serializable
 
 
 @register_keras_serializable("deel-lip", "WeightClipConstraint")
@@ -49,8 +48,8 @@ def __init__(self, scale=1):
         self.scale = scale
 
     def __call__(self, w):
-        c = 1 / (tf.sqrt(tf.cast(tf.size(w), dtype=w.dtype)) * self.scale)
-        return tf.clip_by_value(w, -c, c)
+        c = 1 / (K.sqrt(K.cast(K.size(w), dtype=w.dtype)) * self.scale)
+        return K.clip(w, -c, c)
 
     def get_config(self):
         return {"scale": self.scale}
@@ -67,7 +66,7 @@ def __init__(self, eps=1e-7):
         self.eps = eps
 
     def __call__(self, w):
-        return w / (tf.sqrt(tf.reduce_sum(tf.square(w), keepdims=False)) + self.eps)
+        return w / (K.sqrt(K.sum(K.square(w), keepdims=False)) + self.eps)
 
     def get_config(self):
         return {"eps": self.eps}
@@ -95,15 +94,15 @@ def __init__(
             eps_spectral (float): stopping criterion for the iterative power algorithm.
             eps_bjorck (float): stopping criterion Bjorck algorithm.
             beta_bjorck (float): beta parameter in bjorck algorithm.
-            u (tf.Tensor): vector used for iterated power method, can be set to None
+            u (Tensor): vector used for iterated power method, can be set to None
                 (used for serialization/deserialization purposes).
         """
         self.eps_spectral = eps_spectral
         self.eps_bjorck = eps_bjorck
         self.beta_bjorck = beta_bjorck
         self.k_coef_lip = k_coef_lip
-        if not (isinstance(u, tf.Tensor) or (u is None)):
-            u = tf.convert_to_tensor(u)
+        if not (K.is_tensor(u) or (u is None)):
+            u = K.convert_to_tensor(u)
         self.u = u
         super(SpectralConstraint, self).__init__()
 

deel/lip/initializers.py

@@ -7,19 +7,19 @@
 matrix initialization.
 They can be used as kernel initializers in any Keras layer.
 """
-from tensorflow.keras.initializers import Initializer
-from tensorflow.keras import initializers
+import keras
+from keras.saving import register_keras_serializable
+
 from .normalizers import (
     reshaped_kernel_orthogonalization,
     DEFAULT_EPS_SPECTRAL,
     DEFAULT_EPS_BJORCK,
     DEFAULT_BETA_BJORCK,
 )
-from tensorflow.keras.utils import register_keras_serializable
 
 
 @register_keras_serializable("deel-lip", "SpectralInitializer")
-class SpectralInitializer(Initializer):
+class SpectralInitializer(keras.Initializer):
     def __init__(
         self,
         eps_spectral=DEFAULT_EPS_SPECTRAL,
@@ -44,7 +44,7 @@ def __init__(
         self.eps_bjorck = eps_bjorck
         self.beta_bjorck = beta_bjorck
         self.k_coef_lip = k_coef_lip
-        self.base_initializer = initializers.get(base_initializer)
+        self.base_initializer = keras.initializers.get(base_initializer)
         super(SpectralInitializer, self).__init__()
 
     def __call__(self, shape, dtype=None, partition_info=None):
@@ -65,5 +65,5 @@ def get_config(self):
             "eps_bjorck": self.eps_bjorck,
             "beta_bjorck": self.beta_bjorck,
             "k_coef_lip": self.k_coef_lip,
-            "base_initializer": initializers.serialize(self.base_initializer),
+            "base_initializer": keras.initializers.serialize(self.base_initializer),
         }