steal from submodules

Author: xadupre

_unittests/ut_helpers/test_torch_test_helper.py

@@ -144,6 +144,39 @@ def forward(self, x, y):
         self.assertEqualAny(restored["main", 0, "O"], res1)
         self.assertEqualAny(restored["main", 0, "O"], res2)
 
+    @hide_stdout()
+    def test_steal_forward_submodules(self):
+        class SubModel(torch.nn.Module):
+            def forward(self, x):
+                return x * x
+
+        class Model(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.s1 = SubModel()
+                self.s2 = SubModel()
+
+            def forward(self, x, y):
+                return self.s1(x) + self.s2(y)
+
+        inputs = torch.rand(3, 4), torch.rand(3, 4)
+        model = Model()
+        dump_file = self.get_dump_file("test_steal_forward_submodules.onnx")
+        with steal_forward(model, submodules=True, dump_file=dump_file):
+            model(*inputs)
+        restored = create_input_tensors_from_onnx_model(dump_file)
+        self.assertEqual(
+            [
+                ("", 0, "I"),
+                ("", 0, "O"),
+                ("s1", 0, "I"),
+                ("s1", 0, "O"),
+                ("s2", 0, "I"),
+                ("s2", 0, "O"),
+            ],
+            sorted(restored),
+        )
+
     def test_replace_string_by_dynamic(self):
         example = {
             "input_ids": {0: "batch_size", 1: "sequence_length"},

onnx_diagnostic/helpers/torch_test_helper.py

@@ -54,6 +54,7 @@ def steal_forward(
     ],
     fprint: Callable = string_type,
     dump_file: Optional[str] = None,
+    submodules: bool = False,
     **kwargs,
 ):
     """
@@ -70,12 +71,25 @@ def steal_forward(
     :param dump_file: dumps stolen inputs and outputs in an onnx model,
         they can be restored with :func:`create_input_tensors_from_onnx_model
         <onnx_diagnostic.helpers.mini_onnx_builder.create_input_tensors_from_onnx_model>`
+    :param submodules: if True and model is a module, the list extended with all the submodules
+        the module contains
     """
+    assert not submodules or isinstance(
+        model, torch.nn.Module
+    ), f"submodules can only be True if model is a module but is is {type(model)}."
     context = dict(iteration=0, **kwargs)
     if "with_shape" not in context and fprint == string_type:
         context["with_shape"] = True
     if not isinstance(model, list):
-        model = [model]
+        if submodules:
+            models = []
+            for idx, m in model.named_modules():
+                level = str(idx).split(".")
+                ll = len(level)
+                models.append((f"{'  ' * ll}{idx}", m))
+            model = models
+        else:
+            model = [model]
     keep_model_forward = {}
     storage: Optional[Dict[Any, Any]] = {} if dump_file else None
     for mt in model: