add self.built=False to counter issue #20871 -corrected in PR #20880- + update parameters usage in InvertibleUpDownSampling

Author: franckma31

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

@@ -93,7 +93,6 @@ def evaluate_lip_const(model: keras.Model, x, eps=1e-4, seed=None):
     return keras.ops.max(spectral_norms).numpy()
 
 
-
 def _padding_circular(x, circular_paddings):
     """Add circular padding to a 4-D tensor. Only channels_last is supported."""
     if circular_paddings is None:

setup.cfg

@@ -19,7 +19,7 @@ deps =
 
 
 commands =
-    python -m unittest
+    pytest tests
 
 [testenv:py{39,310,311}-lint]
 skip_install = true

tests/utils_framework.py

@@ -2,10 +2,8 @@
 import os
 import warnings
 from functools import partial
-import pytest
 import numpy as np
 
-import keras
 import keras.utils as K
 import tensorflow as tf
 
@@ -24,9 +22,6 @@
 # from tensorflow.keras.optimizers import SGD, Adam
 from keras.optimizers import SGD, Adam
 
-# from tensorflow.keras.losses import CategoricalCrossentropy as CategoricalCrossentropy
-from keras.losses import CategoricalCrossentropy as CategoricalCrossentropy
-
 # from tensorflow.keras.layers import Layer as tLayer
 from keras.layers import Layer as tLayer
 
@@ -36,16 +31,6 @@
 # from tensorflow.keras.layers import Flatten
 from keras.layers import Flatten
 
-
-from keras.utils import set_random_seed as set_seed
-
-# from tensorflow.keras import backend as K
-# from tensorflow.keras.metrics import mse as tmse
-from keras.metrics import MeanSquaredError as tmse
-
-# from tensorflow.keras.losses import MeanSquaredError as MeanSquaredError
-from keras.losses import MeanSquaredError as MeanSquaredError
-
 # from tensorflow import int32 as type_int32
 from tensorflow import int32 as type_int32
 
@@ -79,9 +64,6 @@
 # from tensorflow.keras.layers import Concatenate as tConcatenate
 from keras.layers import Concatenate as tConcatenate
 
-import sys
-
-sys.path.append(".")
 from deel.lip.activations import GroupSort as GroupSort
 from deel.lip.activations import GroupSort2 as GroupSort2
 from deel.lip.activations import Householder as HouseHolder
@@ -118,7 +100,6 @@
 from deel.lip.losses import MulticlassHinge as HingeMulticlassLoss
 from deel.lip.losses import MulticlassHKR as HKRMulticlassLoss
 
-####from deel.lip.losses import MulticlassSoftHKR as SoftHKRMulticlassLoss
 from deel.lip.losses import MultiMargin as MultiMarginLoss
 from deel.lip.losses import TauCategoricalCrossentropy as TauCategoricalCrossentropyLoss
 from deel.lip.losses import (
@@ -176,7 +157,7 @@
     "CondenseCallback",
     "MonitorCallback",
     "Sequential",
-    "ScaledGlobalL2NormPool2d",
+    "ScaledAdaptativeL2NormPool2d",
     "evaluate_lip_const",
     "DEFAULT_NITER_SPECTRAL_INIT",
     "Loss",
@@ -209,9 +190,7 @@
 ]
 
 FIT = "fit"
-# if tf.__version__.startswith("2.0") else "fit"
 EVALUATE = "evaluate"
-# if tf.__version__.startswith("2.0") else "evaluate"
 
 MODEL_PATH = "model"
 EXTENSION = ".keras"