Pytest unit tests, new losses support, and normalization enhancement

deel/torchlip/modules/module.py

@@ -31,6 +31,7 @@
 import abc
 import copy
 import logging
+import warnings
 import math
 from collections import OrderedDict
 from typing import Any
@@ -40,8 +41,27 @@
 import torch.nn as nn
 import torch.nn.utils.parametrize as parametrize
 from torch.nn import Sequential as TorchSequential
+from torch import reshape
 
-logger = logging.getLogger("deel.torchlip")
+
+def _is_supported_1lip_layer(layer):
+    """Return True if the Keras layer is 1-Lipschitz. Note that in some cases, the layer
+    is 1-Lipschitz for specific set of parameters.
+    """
+    supported_1lip_layers = (
+        torch.nn.Softmax,
+        torch.nn.Flatten,
+        torch.nn.Identity,
+        torch.nn.ReLU,
+        torch.nn.Sigmoid,
+        torch.nn.Tanh,
+        Reshape,
+    )
+    if isinstance(layer, supported_1lip_layers):
+        return True
+    elif isinstance(layer, torch.nn.MaxPool2d):
+        return layer.kernel_size <= layer.stride
+    return False
 
 
 def vanilla_model(model: nn.Module):
@@ -133,13 +153,13 @@ def __init__(
 
         # Force the Lipschitz coefficient:
         n_layers = np.sum(
-            (isinstance(layer, LipschitzModule) for layer in self.children())
+            [isinstance(layer, LipschitzModule) for layer in self.children()]
         )
         for module in self.children():
             if isinstance(module, LipschitzModule):
                 module._coefficient_lip = math.pow(k_coef_lip, 1 / n_layers)
-            else:
-                logger.warning(
+            elif _is_supported_1lip_layer(module) is not True:
+                warnings.warn(
                     "Sequential model contains a layer which is not a Lipschitz layer: {}".format(  # noqa: E501
                         module
                     )
@@ -163,3 +183,12 @@ def vanilla_export(self):
             else:
                 layers.append((name, copy.deepcopy(layer)))
         return TorchSequential(OrderedDict(layers))
+
+
+class Reshape(torch.nn.Module):
+    def __init__(self, target_shape):
+        super(Reshape, self).__init__()
+        self.target_shape = target_shape
+
+    def forward(self, x):
+        return reshape(x, self.target_shape)

tests/test_models.py

@@ -290,7 +290,7 @@ def test_warning_unsupported_1Lip_layers():
     # Check that unsupported layers raise a warning
     unsupported_layers = [
         uft.get_instance_framework(
-            tMaxPool2d, {"pool_size": 3, "strides": 2}
+            tMaxPool2d, {"kernel_size": 3, "stride": 2}
         ),  # kl.MaxPool2d(),
         uft.get_instance_framework(tAdd, {}),  # kl.Add(),
         uft.get_instance_framework(tConcatenate, {}),  # kl.Concatenate(),
@@ -309,8 +309,8 @@ def test_warning_unsupported_1Lip_layers():
     #         unsupported_layers.append(kl.Activation("gelu"))
 
     for lay in unsupported_layers:
-        with pytest.warns(Warning):
-            if lay is not None:
+        if lay is not None:
+            with pytest.warns(Warning):
                 _ = uft.generate_k_lip_model(
                     Sequential,
                     {"layers": [lay]},