Better SVD computation for convolution

.github/workflows/python-linters.yml

@@ -16,9 +16,9 @@ jobs:
         python-version: [3.7, "3.10"]
 
     steps:
-    - uses: actions/checkout@v3
+    - uses: actions/checkout@v4
     - name: Set up Python ${{ matrix.python-version }}
-      uses: actions/setup-python@v4
+      uses: actions/setup-python@v5
       with:
         python-version: ${{ matrix.python-version }}
     - name: Install dependencies

.github/workflows/python-publish.yml

@@ -13,9 +13,9 @@ jobs:
     runs-on: ubuntu-latest
 
     steps:
-    - uses: actions/checkout@v3
+    - uses: actions/checkout@v4
     - name: Set up Python
-      uses: actions/setup-python@v4
+      uses: actions/setup-python@v5
       with:
         python-version: '3.x'
     - name: Install dependencies

.github/workflows/python-tests.yml

@@ -24,9 +24,9 @@ jobs:
             tf-version: latest
 
     steps:
-    - uses: actions/checkout@v3
+    - uses: actions/checkout@v4
     - name: Set up Python ${{ matrix.python-version }}
-      uses: actions/setup-python@v4
+      uses: actions/setup-python@v5
       with:
         python-version: ${{ matrix.python-version }}
     - name: Install dependencies

deel/lip/compute_layer_sv.py

@@ -15,6 +15,8 @@
 """
 
 import numpy as np
+import scipy.sparse as sps
+import scipy.sparse.linalg as spsl  # Required for scipy<=1.7.3
 import tensorflow as tf
 
 from .layers import Condensable, GroupSort, MaxMin
@@ -53,28 +55,25 @@ def _generate_conv_matrix(layer, input_sizes):
     Returns:
         np.array: the equivalent matrix of the convolutional layer.
     """
+    # Clone layer (as a model) and remove bias and activation
     single_layer_model = tf.keras.models.Sequential(
         [tf.keras.layers.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]
-    w_eqmatrix = None
-    start_index = 0
-    for ch in range(channel_in):
+    single_layer_model = tf.keras.models.clone_model(single_layer_model)
+    single_layer_model.set_weights(layer.get_weights())
+    single_layer_model.layers[0].use_bias = False
+    single_layer_model.layers[0].activation = None
+
+    dirac_inp = np.zeros((input_sizes[2],) + input_sizes[1:], dtype=np.float32)
+    stack = []
+    for ch in range(input_sizes[-1]):
         for ii in range(input_sizes[1]):
             dirac_inp[:, ii, :, ch] = np.eye(input_sizes[2])
             out_pred = single_layer_model(dirac_inp)
-            if w_eqmatrix is None:
-                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)
-            )
+            out_pred = tf.reshape(out_pred, (input_sizes[2], -1))
+            stack.append(sps.coo_matrix(out_pred.numpy()))
             dirac_inp = 0.0 * dirac_inp
-            start_index += input_sizes[2]
-    return w_eqmatrix
+    return sps.vstack(stack)
 
 
 def _compute_sv_conv2d_layer(layer, input_sizes):
@@ -94,8 +93,9 @@ def _compute_sv_conv2d_layer(layer, input_sizes):
         tuple: min and max singular values
     """
     w_eqmatrix = _generate_conv_matrix(layer, input_sizes)
-    svd = np.linalg.svd(w_eqmatrix, compute_uv=False)
-    return (np.min(svd), np.max(svd))
+    svd_max = spsl.svds(w_eqmatrix, k=1, return_singular_vectors=False)
+    svd_min = spsl.svds(w_eqmatrix, k=1, which="SM", return_singular_vectors=False)
+    return (svd_min[0], svd_max[0])
 
 
 def _compute_sv_activation(layer, input_sizes=None):

setup.py

@@ -49,7 +49,7 @@
     url="https://github.com/deel-ai/deel-lip",
     packages=setuptools.find_namespace_packages(include=["deel.*"]),
     include_package_data=True,
-    install_requires=["numpy", "tensorflow~=2.2"],
+    install_requires=["numpy", "scipy", "tensorflow~=2.2"],
     license="MIT",
     extras_require={"dev": dev_requires, "docs": docs_requires},
     classifiers=[

tests/test_compute_layer_sv.py

@@ -343,6 +343,23 @@ def test_spectralconv2d(self):
                     test_SVmin=False,
                     callbacks=[],
                 ),
+                dict(  # SpectralConv2D with bias and activation
+                    layer_type=SpectralConv2D,
+                    layer_params={
+                        "filters": 2,
+                        "kernel_size": (3, 3),
+                        "activation": "relu",
+                    },
+                    batch_size=100,
+                    steps_per_epoch=125,
+                    epochs=5,
+                    input_shape=(5, 5, 1),
+                    k_lip_data=1.0,
+                    k_lip_model=1.0,
+                    k_lip_tolerance_factor=1.02,
+                    test_SVmin=False,
+                    callbacks=[],
+                ),
                 dict(
                     layer_type=SpectralConv2D,
                     layer_params={