Merge pull request #16 from HeinrichAD/feature/keep-layer-names-flag

Keep the layer names when vanilla exporting

deel/torchlip/modules/module.py

@@ -29,6 +29,7 @@
 for condensation and vanilla exportation.
 """
 import abc
+from collections import OrderedDict
 import copy
 import logging
 import math
@@ -118,9 +119,9 @@ def vanilla_export(self):
             A Vanilla torch.nn.Sequential model.
         """
         layers = []
-        for layer in self.children():
+        for name, layer in self.named_children():
             if isinstance(layer, LipschitzModule):
-                layers.append(layer.vanilla_export())
+                layers.append((name, layer.vanilla_export()))
             else:
-                layers.append(copy.deepcopy(layer))
-        return TorchSequential(*layers)
+                layers.append((name, copy.deepcopy(layer)))
+        return TorchSequential(OrderedDict(layers))

tests/test_vanilla_export.py

@@ -0,0 +1,59 @@
+# -*- coding: utf-8 -*-
+# Copyright IRT Antoine de Saint Exupéry et Université Paul Sabatier Toulouse III - All
+# rights reserved. DEEL is a research program operated by IVADO, IRT Saint Exupéry,
+# CRIAQ and ANITI - https://www.deel.ai/
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+# Copyright IRT Antoine de Saint Exupéry et Université Paul Sabatier Toulouse III - All
+# rights reserved. DEEL is a research program operated by IVADO, IRT Saint Exupéry,
+# CRIAQ and ANITI - https://www.deel.ai/
+# =====================================================================================
+from deel import torchlip
+
+from collections import OrderedDict
+
+
+def get_named_model():
+    return torchlip.Sequential(
+        OrderedDict(
+            [
+                (
+                    "features",
+                    torchlip.Sequential(
+                        torchlip.SpectralLinear(2, 256),
+                        torchlip.FullSort(),
+                        torchlip.SpectralLinear(256, 128),
+                        torchlip.FullSort(),
+                        torchlip.SpectralLinear(128, 64),
+                        torchlip.FullSort(),
+                    ),
+                ),
+                (
+                    "classifier",
+                    torchlip.FrobeniusLinear(64, 1),
+                ),
+            ]
+        )
+    )
+
+
+def test_vanilla_export_with_named_layers():
+    model = get_named_model()
+    names = [name for name, _ in model.vanilla_export().named_children()]
+    assert names == ["features", "classifier"]