add model_statistics

Author: xadupre

_unittests/ut_helpers/test_helper.py

@@ -127,7 +127,7 @@ def test_print_pretty_onnx(self):
         )
         self.print_onnx(proto)
         self.print_model(proto)
-        self.dump_onnx("test_print_pretty_onnx", proto)
+        self.dump_onnx("test_print_pretty.onnx", proto)
         self.check_ort(proto)
         self.assertNotEmpty(proto)
         self.assertEmpty(None)

_unittests/ut_helpers/test_torch_test_helper.py

@@ -9,6 +9,7 @@
     dummy_llm,
     to_numpy,
     is_torchdynamo_exporting,
+    model_statistics,
     steal_forward,
     replace_string_by_dynamic,
     to_any,
@@ -279,6 +280,25 @@ def test_torch_deepcopy_sliding_windon_cache(self):
     def test_torch_deepcopy_none(self):
         self.assertEmpty(torch_deepcopy(None))
 
+    def test_model_statistics(self):
+        class Model(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.p1 = torch.nn.Parameter(torch.tensor([1], dtype=torch.float32))
+                self.b1 = torch.nn.Buffer(torch.tensor([1], dtype=torch.float32))
+
+            def forward(self, x, y=None):
+                return x + y + self.p1 + self.b1
+
+        model = Model()
+        x, y = torch.rand((3, 4)), torch.rand((3, 4))
+        model(x, y)
+        stat = model_statistics(model)
+        self.assertEqual(
+            {"type": "Model", "n_modules": 1, "param_size": 4, "buffer_size": 4, "float32": 8},
+            stat,
+        )
+
 
 if __name__ == "__main__":
     unittest.main(verbosity=2)

onnx_diagnostic/helpers/torch_test_helper.py

@@ -487,3 +487,35 @@ def torch_deepcopy(value: Any) -> Any:
     # We should have a code using serialization, deserialization assuming a model
     # cannot be exported without them.
     raise NotImplementedError(f"torch_deepcopy not implemented for type {type(value)}")
+
+
+def model_statistics(model: torch.nn.Module):
+    """Returns statistics on a model in a dictionary."""
+    n_subs = len(list(model.modules()))
+    sizes = {}
+    param_size = 0
+    for param in model.parameters():
+        size = param.nelement() * param.element_size()
+        param_size += size
+        name = str(param.dtype).replace("torch.", "")
+        if name not in sizes:
+            sizes[name] = 0
+        sizes[name] += size
+
+    buffer_size = 0
+    for buffer in model.buffers():
+        size = buffer.nelement() * buffer.element_size()
+        buffer_size += size
+        name = str(buffer.dtype).replace("torch.", "")
+        if name not in sizes:
+            sizes[name] = 0
+        sizes[name] += size
+
+    res = dict(
+        type=model.__class__.__name__,
+        n_modules=n_subs,
+        param_size=size,
+        buffer_size=buffer_size,
+    )
+    res.update(sizes)
+    return res