feat(compute_layer_sv): sparse SVD computation

deel/lip/compute_layer_sv.py

@@ -15,6 +15,7 @@
 """
 
 import numpy as np
+import scipy.sparse as sps
 import tensorflow as tf
 
 from .layers import Condensable, GroupSort, MaxMin
@@ -62,25 +63,16 @@ def _generate_conv_matrix(layer, input_sizes):
     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:])  # 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):
+    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):
@@ -100,8 +92,11 @@ 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 = sps.linalg.svds(w_eqmatrix, k=1, return_singular_vectors=False)
+    svd_min = sps.linalg.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=[