Executorch Model Size Analysis Tool (#68)

Summary:
Pull Request resolved: #68

Adds a Size Analysis Tool for Executorch models, which provides information about tensor data and delegate blobs in a model. It can be used by calling ```generate_model_size_information``` directly, or from the command line
```
buck2 run executorch/sdk/size_analysis_tool:size_analysis_tool -- --etrecord_path [...] --output_path [...]
```

Reviewed By: mcr229

Differential Revision: D48255922

fbshipit-source-id: 2d68bcb86ea347d82435bce095bc0669d3e889b5

Author: salilsdesai

sdk/size_analysis_tool/TARGETS

@@ -0,0 +1,49 @@
+load("@fbcode_macros//build_defs:python_binary.bzl", "python_binary")
+load("@fbcode_macros//build_defs:python_library.bzl", "python_library")
+load("@fbcode_macros//build_defs:python_unittest.bzl", "python_unittest")
+
+python_library(
+    name = "size_analysis_tool_lib",
+    srcs = [
+        "size_analysis_tool.py",
+    ],
+    visibility = ["PUBLIC"],
+    deps = [
+        "//caffe2:torch",
+        "//executorch/exir:lib",
+        "//executorch/exir/backend:backend_api",
+        "//executorch/sdk/etrecord:etrecord",
+    ],
+)
+
+python_binary(
+    name = "size_analysis_tool",
+    srcs = [
+        "size_analysis_tool.py",
+    ],
+    main_module = "executorch.sdk.size_analysis_tool.size_analysis_tool",
+    visibility = ["PUBLIC"],
+    deps = [
+        "//caffe2:torch",
+        "//executorch/exir:lib",
+        "//executorch/exir/backend:backend_api",
+        "//executorch/sdk/etrecord:etrecord",
+    ],
+)
+
+python_unittest(
+    name = "size_analysis_tool_test",
+    srcs = [
+        "size_analysis_tool.py",
+        "size_analysis_tool_test.py",
+    ],
+    deps = [
+        "//caffe2:torch",
+        "//executorch/backends/xnnpack/partition:xnnpack_partitioner",
+        "//executorch/backends/xnnpack/utils:xnnpack_utils",
+        "//executorch/exir:lib",
+        "//executorch/exir/backend:backend_api",
+        "//executorch/exir/passes:spec_prop_pass",
+        "//executorch/sdk/etrecord:etrecord",
+    ],
+)

sdk/size_analysis_tool/size_analysis_tool.py

@@ -0,0 +1,185 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+import argparse
+import json
+from typing import Any, Callable, Dict, List, Optional, Tuple
+
+import torch
+
+from executorch.exir import ExportedProgram
+from executorch.exir.backend.backend_api import LoweredBackendModule
+from executorch.sdk.etrecord import parse_etrecord
+from executorch.sdk.etrecord._etrecord import ETRecordReservedFileNames
+
+
+def _get_tensor_data(node: torch.fx.Node, tensor: torch.Tensor) -> Dict[str, Any]:
+    return {
+        "name": node.name,
+        "numel": tensor.numel(),
+        "dtype": str(tensor.dtype)[6:],  # Remove "torch." prefix
+        "element_size": tensor.element_size(),
+        "shape": list(tensor.shape),
+        "num_bytes": tensor.element_size() * tensor.numel(),
+        "nn_module_stack": (
+            str(node.meta["nn_module_stack"])
+            if "nn_module_stack" in node.meta
+            else None
+        ),
+    }
+
+
+def _get_delegate_blob_data(
+    node: torch.fx.Node,
+    lowered_backend_module: LoweredBackendModule,
+    delegate_deserializers: Optional[
+        Dict[str, Callable[[bytes], Dict[str, Any]]]
+    ] = None,
+) -> Dict[str, Any]:
+    delegate_blob_data = {
+        "name": node.name,
+        "backend_id": lowered_backend_module.backend_id,
+        "num_bytes": len(lowered_backend_module.processed_bytes),
+    }
+    if (
+        delegate_deserializers is not None
+        and lowered_backend_module.backend_id in delegate_deserializers
+    ):
+        delegate_blob_data.update(
+            delegate_deserializers[lowered_backend_module.backend_id](
+                lowered_backend_module.processed_bytes
+            )
+        )
+
+    return delegate_blob_data
+
+
+def _get_nested_model_data(
+    graph_module: torch.fx.GraphModule,
+    delegate_deserializers: Optional[
+        Dict[str, Callable[[bytes], Dict[str, Any]]]
+    ] = None,
+    tensor_data: Optional[List[Dict[str, Any]]] = None,
+    delegate_blob_data: Optional[List[Dict[str, Any]]] = None,
+) -> Tuple[List[Dict[str, Any]], List[Dict[str, Any]]]:
+    if tensor_data is None:
+        tensor_data = []
+
+    if delegate_blob_data is None:
+        delegate_blob_data = []
+
+    for node in graph_module.graph.nodes:
+        if node.op == "get_attr":
+            node_attr = getattr(node.graph.owning_module, node.target)
+            if isinstance(node_attr, torch.Tensor):
+                tensor_data.append(_get_tensor_data(node, node_attr))
+            elif isinstance(node_attr, torch.fx.GraphModule):
+                _get_nested_model_data(
+                    node_attr, delegate_deserializers, tensor_data, delegate_blob_data
+                )
+            elif isinstance(node_attr, LoweredBackendModule):
+                delegate_blob_data.append(
+                    _get_delegate_blob_data(node, node_attr, delegate_deserializers)
+                )
+
+    return (tensor_data, delegate_blob_data)
+
+
+def generate_model_size_information(
+    model: ExportedProgram,
+    delegate_deserializers: Optional[
+        Dict[str, Callable[[bytes], Dict[str, Any]]]
+    ] = None,
+    flatbuffer: Optional[bytes] = None,
+) -> Dict[str, Any]:
+    """
+    Generate a json-serializable Dict containing information about a model's
+    size. This includes data about individual tensors and delegate blobs.
+    Optionally:
+    - delegate_deserializers can be provided to manually specify additional
+      information to include for delegate blobs for specific backends.
+    - flatbuffer can be provided to include a comparison of total tensor data
+      size to overall model size
+    """
+
+    tensor_and_delegate_blob_data = _get_nested_model_data(
+        model.graph_module, delegate_deserializers
+    )
+
+    for data_list in tensor_and_delegate_blob_data:
+        data_list.sort(key=lambda data: data["num_bytes"], reverse=True)
+
+    (tensor_data, delegate_blob_data) = tensor_and_delegate_blob_data
+
+    total_tensor_data_size = sum(data["num_bytes"] for data in tensor_data)
+    total_delegate_blob_data_size = sum(
+        data["num_bytes"] for data in delegate_blob_data
+    )
+    overview = {
+        "total_tensor_data_size": total_tensor_data_size,
+        "total_delegate_blob_data_size": total_delegate_blob_data_size,
+    }
+    if flatbuffer is not None:
+        model_size = len(flatbuffer)
+        overview.update(
+            {
+                "serialization_metadata_size": (
+                    model_size - total_tensor_data_size - total_delegate_blob_data_size
+                ),
+                "model_size": model_size,
+            }
+        )
+
+    return {
+        "tensor_data": tensor_data,
+        "delegate_blob_data": delegate_blob_data,
+        "overview": overview,
+    }
+
+
+def parse_args():
+    parser = argparse.ArgumentParser()
+
+    parser.add_argument(
+        "--etrecord_path",
+        required=True,
+        help="The path to the ETRecord for the model to generate size information for",
+    )
+
+    parser.add_argument(
+        "--output_path",
+        default="model_size_information.json",
+        help="The output path for the model size information as a json file",
+    )
+
+    args = parser.parse_args()
+    return args
+
+
+def main():
+    args = parse_args()
+
+    etrecord = parse_etrecord(args.etrecord_path)
+
+    all_model_size_information = [
+        generate_model_size_information(
+            model=exported_program,
+            delegate_deserializers=None,
+            flatbuffer=(
+                etrecord.program_buffer
+                if name == ETRecordReservedFileNames.ET_DIALECT_GRAPH_MODULE
+                else None
+            ),
+        )
+        for (name, exported_program) in etrecord.graph_map.items()
+    ]
+
+    with open(args.output_path, "w") as f:
+        f.write(json.dumps(all_model_size_information))
+
+
+if __name__ == "__main__":
+    main()

sdk/size_analysis_tool/size_analysis_tool_test.py

@@ -0,0 +1,115 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+import unittest
+
+import torch
+from executorch.backends.xnnpack.partition.xnnpack_partitioner import (
+    XnnpackFloatingPointPartitioner,
+)
+from executorch.backends.xnnpack.utils.configs import (
+    get_xnnpack_executorch_backend_config,
+)
+from executorch.backends.xnnpack.utils.utils import capture_graph_for_xnnpack
+from executorch.exir.backend.backend_api import to_backend, validation_disabled
+from executorch.exir.passes.spec_prop_pass import SpecPropPass
+
+from executorch.sdk.size_analysis_tool.size_analysis_tool import (
+    generate_model_size_information,
+)
+
+
+class SizeAnalysisToolTest(unittest.TestCase):
+    def test_generate_model_size_analysis(self):
+        class MyModel(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.sigmoid = torch.nn.Sigmoid()
+                self.conv3d = torch.nn.Conv3d(
+                    in_channels=4, out_channels=2, kernel_size=3
+                )
+                self.conv2d = torch.nn.Conv2d(
+                    in_channels=5,
+                    out_channels=2,
+                    kernel_size=3,
+                )
+                self.conv_transpose2d = torch.nn.ConvTranspose2d(
+                    in_channels=2, out_channels=4, kernel_size=2
+                )
+
+            def forward(self, x):
+                x = self.sigmoid(x)
+                x = self.conv3d(x)
+                x = self.conv2d(x)
+                x = self.conv_transpose2d(x)
+                return x
+
+        mm = MyModel()
+        mm.eval()
+
+        test_input = torch.ones(size=(4, 7, 5, 6), dtype=torch.float)
+
+        edge_program = capture_graph_for_xnnpack(mm, (test_input,))
+        partitioner = XnnpackFloatingPointPartitioner
+
+        with validation_disabled():
+            delegated_program = edge_program
+            delegated_program.exported_program = to_backend(
+                edge_program.exported_program, partitioner
+            )
+
+        program = delegated_program.to_executorch(
+            get_xnnpack_executorch_backend_config([SpecPropPass()]),
+        )
+
+        size_information = generate_model_size_information(
+            model=program,
+            delegate_deserializers=None,
+            flatbuffer=program.buffer,
+        )
+
+        # Number of Elements -> Other tensor data
+        exepected_tensor_data = {
+            # Conv3d Weight
+            216: {
+                "dtype": "float32",
+                "element_size": 4,
+                "shape": [2, 4, 3, 3, 3],
+                "num_bytes": 864,
+            },
+            # ConvTranspose2d Weight
+            32: {
+                "dtype": "float32",
+                "element_size": 4,
+                "shape": [2, 4, 2, 2],
+                "num_bytes": 128,
+            },
+            # ConvTranspose2d Bias
+            4: {
+                "dtype": "float32",
+                "element_size": 4,
+                "shape": [4],
+                "num_bytes": 16,
+            },
+            # Conv3d Bias
+            2: {
+                "dtype": "float32",
+                "element_size": 4,
+                "shape": [2],
+                "num_bytes": 8,
+            },
+        }
+
+        self.assertEqual(
+            len(size_information["tensor_data"]), len(exepected_tensor_data)
+        )
+
+        for tensor in size_information["tensor_data"]:
+            for (k, v) in exepected_tensor_data[tensor["numel"]].items():
+                self.assertEqual(tensor[k], v)
+
+        # Two delegate blobs: sigmoid and conv2d
+        self.assertEqual(len(size_information["delegate_blob_data"]), 2)