steal_append

Author: xadupre

_unittests/ut_helpers/test_torch_test_helper.py

@@ -10,6 +10,7 @@
     to_numpy,
     is_torchdynamo_exporting,
     model_statistics,
+    steal_append,
     steal_forward,
     replace_string_by_dynamic,
     to_any,
@@ -145,6 +146,36 @@ 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_dump_file_steal_append(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):
+                sx = self.s1(x)
+                steal_append("sx", sx)
+                return sx + self.s2(y)
+
+        inputs = torch.rand(3, 4), torch.rand(3, 4)
+        model = Model()
+        dump_file = self.get_dump_file("test_steal_forward_dump_file.onnx")
+        with steal_forward(model, dump_file=dump_file):
+            model(*inputs)
+            model(*inputs)
+        self.assertExists(dump_file)
+        restored = create_input_tensors_from_onnx_model(dump_file)
+        self.assertEqual(
+            {("", 1, "I"), ("", 1, "O"), "sx", ("", 0, "O"), "sx_1", ("", 0, "I")},
+            set(restored),
+        )
+
     @hide_stdout()
     def test_steal_forward_submodules(self):
         class SubModel(torch.nn.Module):

onnx_diagnostic/helpers/torch_test_helper.py

@@ -47,6 +47,35 @@ def _forward_(
     return res
 
 
+_additional_stolen_objects = {}
+
+
+def steal_append(name: str, obj: Any):
+    """
+    When outside a forward method, it is still possible to add
+    a python object which contains tensors and dump after the execution
+    of the model.
+
+    .. code-block:: python
+
+        steal_append("quantize", [t1, t2])
+
+    The same code can executed multiple times, then
+    the name can extended with a number.
+    """
+    if name in _additional_stolen_objects:
+        i = 1
+        n = f"{name}_{i}"
+        while n in _additional_stolen_objects:
+            i += 1
+            n = f"{name}_{i}"
+        print(f"-- stolen {name!r} renamed in {n!r}: {string_type(obj, with_shape=True)}")
+        _additional_stolen_objects[n] = obj
+    else:
+        print(f"-- stolen {name!r}: {string_type(obj, with_shape=True)}")
+        _additional_stolen_objects[name] = obj
+
+
 @contextlib.contextmanager
 def steal_forward(
     model: Union[
@@ -111,7 +140,11 @@ def forward(self, x, y):
                 print("input", k, args, kwargs)
             else:
                 print("output", k, v)
+
+    Function :func:`steal_append` can be used to dump more tensors.
     """
+    # We clear the cache.
+    _additional_stolen_objects.clear()
     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)}."
@@ -147,6 +180,10 @@ def forward(self, x, y):
         for f in keep_model_forward.values():
             f[0].forward = f[1]
         if dump_file:
+            # Let's add the cached tensor
+            storage.update(_additional_stolen_objects)
+            # We clear the cache.
+            _additional_stolen_objects.clear()
             proto = create_onnx_model_from_input_tensors(storage)
             onnx.save(
                 proto,