use lp_lppool2d in ScaledAdaptativeL2NormPool2d ( #23 (comment)) + name change Global to Adaptative

Author: franckma31

deel/torchlip/modules/__init__.py

@@ -75,7 +75,7 @@
 from .pooling import ScaledAdaptiveAvgPool2d
 from .pooling import ScaledAvgPool2d
 from .pooling import ScaledL2NormPool2d
-from .pooling import ScaledGlobalL2NormPool2d
+from .pooling import ScaledAdaptativeL2NormPool2d
 from .upsampling import InvertibleUpSampling
 from .normalization import LayerCentering
 from .normalization import BatchCentering

deel/torchlip/modules/pooling.py

@@ -29,6 +29,7 @@
 
 import numpy as np
 import torch
+import torch.nn.functional as F
 from torch.nn.common_types import _size_2_t
 
 from ..utils import sqrt_with_gradeps
@@ -193,12 +194,13 @@ def vanilla_export(self):
             return self
 
 
-class ScaledGlobalL2NormPool2d(torch.nn.AdaptiveAvgPool2d, LipschitzModule):
+class ScaledAdaptativeL2NormPool2d(
+    torch.nn.modules.pooling._AdaptiveAvgPoolNd, LipschitzModule
+):
     def __init__(
         self,
         output_size: _size_2_t = (1, 1),
         k_coef_lip: float = 1.0,
-        eps_grad_sqrt: float = 1e-6,
     ):
         """
         Average pooling operation for spatial data, with a lipschitz bound. This
@@ -210,37 +212,24 @@ def __init__(
         Arguments:
             output_size: the target output size has to be (1,1)
             k_coef_lip: the lipschitz factor to ensure
-            eps_grad_sqrt: Epsilon value to avoid numerical instability
-                due to non-defined gradient at 0 in the sqrt function
 
         Input shape:
             4D tensor with shape `(batch_size, channels, rows, cols)`.
 
         Output shape:
             4D tensor with shape `(batch_size, channels, 1, 1)`.
         """
-        if eps_grad_sqrt < 0.0:
-            raise RuntimeError("eps_grad_sqrt must be positive")
         if not isinstance(output_size, tuple) or len(output_size) != 2:
             raise RuntimeError("output_size must be a tuple of 2 integers")
         else:
             if output_size[0] != 1 or output_size[1] != 1:
                 raise RuntimeError("output_size must be (1, 1)")
-        torch.nn.AdaptiveAvgPool2d.__init__(self, output_size)
+        torch.nn.modules.pooling._AdaptiveAvgPoolNd.__init__(self, output_size)
         LipschitzModule.__init__(self, k_coef_lip)
-        self.eps_grad_sqrt = eps_grad_sqrt
 
     def forward(self, input: torch.Tensor) -> torch.Tensor:
-        # coeff = computePoolScalingFactor(input.shape[-2:]) * self._coefficient_lip
-        # avg = torch.nn.AdaptiveAvgPool2d.forward(self, torch.square(input))
-        # return  sqrt_with_gradeps(avg,self.eps_grad_sqrt)* coeff
-        return (  # type: ignore
-            sqrt_with_gradeps(
-                torch.square(input).sum(axis=(2, 3), keepdim=True),
-                self.eps_grad_sqrt,
-            )
-            * self._coefficient_lip
-        )
+        return F.lp_pool2d(input, 2, input.shape[-2:]) * self._coefficient_lip
 
     def vanilla_export(self):
         return self
+

docs/source/basic_example.rst

@@ -54,7 +54,7 @@ The following table indicates which module are safe to use in a Lipschitz networ
      -
    * - :class:`torch.nn.AvgPool2d`\ :raw-html-m2r:`<br>`\ :class:`torch.nn.AdaptiveAvgPool2d`
      - no
-     - :class:`.ScaledAvgPool2d`\ :raw-html-m2r:`<br>`\ :class:`.ScaledAdaptiveAvgPool2d` \ :raw-html-m2r:`<br>` \ :class:`.ScaledL2NormPool2d` \ :raw-html-m2r:`<br>` \ :class:`.ScaledGlobalL2NormPool2d`
+     - :class:`.ScaledAvgPool2d`\ :raw-html-m2r:`<br>`\ :class:`.ScaledAdaptiveAvgPool2d` \ :raw-html-m2r:`<br>` \ :class:`.ScaledL2NormPool2d` \ :raw-html-m2r:`<br>` \ :class:`.ScaledAdaptativeL2NormPool2d`
      - The Lipschitz constant is bounded by ``sqrt(pool_h * pool_w)``.
    * - :class:`Flatten`
      - yes

docs/source/deel.torchlip.rst

@@ -39,7 +39,7 @@ Pooling Layers
 .. autoclass:: ScaledAdaptiveAvgPool2d
 .. autoclass:: ScaledAvgPool2d
 .. autoclass:: ScaledL2NormPool2d
-.. autoclass:: ScaledGlobalL2NormPool2d
+.. autoclass:: ScaledAdaptativeL2NormPool2d
 .. autoclass:: InvertibleDownSampling
 .. autoclass:: InvertibleUpSampling
 

tests/test_layers.py

@@ -44,7 +44,7 @@
     ScaledL2NormPool2d,
     InvertibleDownSampling,
     InvertibleUpSampling,
-    ScaledGlobalL2NormPool2d,
+    ScaledAdaptativeL2NormPool2d,
     Flatten,
     Sequential,
 )
@@ -877,7 +877,7 @@ def test_scaledl2normPool2d(test_params):
 
 
 @pytest.mark.skipif(
-    hasattr(ScaledGlobalL2NormPool2d, "unavailable_class"),
+    hasattr(ScaledAdaptativeL2NormPool2d, "unavailable_class"),
     reason="compute_layer_sv not available",
 )
 @pytest.mark.parametrize(
@@ -890,7 +890,7 @@ def test_scaledl2normPool2d(test_params):
                 "layers": [
                     tInput(uft.to_framework_channel((1, 5, 5))),
                     uft.get_instance_framework(
-                        ScaledGlobalL2NormPool2d, {"data_format": "channels_last"}
+                        ScaledAdaptativeL2NormPool2d, {"data_format": "channels_last"}
                     ),
                 ]
             },
@@ -908,7 +908,7 @@ def test_scaledl2normPool2d(test_params):
                 "layers": [
                     tInput(uft.to_framework_channel((1, 5, 5))),
                     uft.get_instance_framework(
-                        ScaledGlobalL2NormPool2d, {"data_format": "channels_last"}
+                        ScaledAdaptativeL2NormPool2d, {"data_format": "channels_last"}
                     ),
                 ]
             },
@@ -926,7 +926,7 @@ def test_scaledl2normPool2d(test_params):
                 "layers": [
                     tInput(uft.to_framework_channel((1, 5, 5))),
                     uft.get_instance_framework(
-                        ScaledGlobalL2NormPool2d, {"data_format": "channels_last"}
+                        ScaledAdaptativeL2NormPool2d, {"data_format": "channels_last"}
                     ),
                 ]
             },

tests/test_pooling.py

@@ -35,7 +35,7 @@
     ScaledAvgPool2d,
     ScaledAdaptiveAvgPool2d,
     ScaledL2NormPool2d,
-    ScaledGlobalL2NormPool2d,
+    ScaledAdaptativeL2NormPool2d,
 )
 
 
@@ -92,7 +92,7 @@ def test_pooling_simple(layer_type, layer_params):
     "layer_type",
     [
         ScaledAdaptiveAvgPool2d,
-        ScaledGlobalL2NormPool2d,
+        ScaledAdaptativeL2NormPool2d,
     ],
 )
 @pytest.mark.parametrize(
@@ -137,8 +137,8 @@ def test_pooling_global(layer_type, layer_params):
         (ScaledL2NormPool2d, {"kernel_size": (2, 2), "stride": (2, 2)}),
         (ScaledAdaptiveAvgPool2d, {"output_size": (1, 1)}),
         (ScaledAdaptiveAvgPool2d, {"output_size": (1, 1), "k_coef_lip": 2.5}),
-        (ScaledGlobalL2NormPool2d, {"output_size": (1, 1)}),
-        (ScaledGlobalL2NormPool2d, {"output_size": (1, 1), "k_coef_lip": 2.5}),
+        (ScaledAdaptativeL2NormPool2d, {"output_size": (1, 1)}),
+        (ScaledAdaptativeL2NormPool2d, {"output_size": (1, 1), "k_coef_lip": 2.5}),
     ],
 )
 def test_pool_vanilla_export(layer_type, layer_params):
@@ -206,7 +206,7 @@ def test_pool_vanilla_export(layer_type, layer_params):
             [[[[40.0 / math.sqrt(4.0 * 4.0)]], [[18.0 / math.sqrt(4.0 * 4.0)]]]],
         ),
         (
-            ScaledGlobalL2NormPool2d,
+            ScaledAdaptativeL2NormPool2d,
             {"output_size": (1, 1)},
             [[[[math.sqrt(120.0)]], [[math.sqrt(320.0)]]]],
         ),

tests/utils_framework.py

@@ -42,7 +42,7 @@
 from deel.torchlip.modules import ScaledAvgPool2d
 from deel.torchlip.modules import ScaledAdaptiveAvgPool2d
 from deel.torchlip.modules import ScaledL2NormPool2d
-from deel.torchlip.modules import ScaledGlobalL2NormPool2d
+from deel.torchlip.modules import ScaledAdaptativeL2NormPool2d
 from deel.torchlip.modules import InvertibleDownSampling
 from deel.torchlip.modules import InvertibleUpSampling
 from deel.torchlip.modules import LayerCentering
@@ -231,7 +231,7 @@ def get_instance_withcheck(
     ScaledL2NormPool2d: partial(
         get_instance_withreplacement, dict_keys_replace={"data_format": None}
     ),
-    ScaledGlobalL2NormPool2d: partial(
+    ScaledAdaptativeL2NormPool2d: partial(
         get_instance_withreplacement, dict_keys_replace={"data_format": None}
     ),
     SpectralConv2d: partial(