diff --git a/sdk/size_analysis_tool/TARGETS b/sdk/size_analysis_tool/TARGETS new file mode 100644 index 00000000000..c1fd9530fbb --- /dev/null +++ b/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", + ], +) diff --git a/sdk/size_analysis_tool/size_analysis_tool.py b/sdk/size_analysis_tool/size_analysis_tool.py new file mode 100644 index 00000000000..345a6899b18 --- /dev/null +++ b/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() diff --git a/sdk/size_analysis_tool/size_analysis_tool_test.py b/sdk/size_analysis_tool/size_analysis_tool_test.py new file mode 100644 index 00000000000..a8bc53750c9 --- /dev/null +++ b/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)