• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
1# Copyright (c) Meta Platforms, Inc. and affiliates.
2# All rights reserved.
3#
4# This source code is licensed under the BSD-style license found in the
5# LICENSE file in the root directory of this source tree.
6
7import unittest
8
9import torch
10from executorch.backends.xnnpack.partition.xnnpack_partitioner import (
11    XnnpackFloatingPointPartitioner,
12)
13from executorch.backends.xnnpack.utils.configs import (
14    get_xnnpack_executorch_backend_config,
15)
16from executorch.backends.xnnpack.utils.utils import capture_graph_for_xnnpack
17
18from executorch.devtools.size_analysis_tool.size_analysis_tool import (
19    generate_model_size_information,
20)
21from executorch.exir.backend.backend_api import to_backend, validation_disabled
22from executorch.exir.passes.spec_prop_pass import SpecPropPass
23
24
25class SizeAnalysisToolTest(unittest.TestCase):
26    def test_generate_model_size_analysis(self):
27        class MyModel(torch.nn.Module):
28            def __init__(self):
29                super().__init__()
30                self.sigmoid = torch.nn.Sigmoid()
31                self.conv3d = torch.nn.Conv3d(
32                    in_channels=4, out_channels=2, kernel_size=3
33                )
34                self.conv2d = torch.nn.Conv2d(
35                    in_channels=5,
36                    out_channels=2,
37                    kernel_size=3,
38                )
39                self.conv_transpose2d = torch.nn.ConvTranspose2d(
40                    in_channels=2, out_channels=4, kernel_size=2
41                )
42
43            def forward(self, x):
44                x = self.sigmoid(x)
45                x = self.conv3d(x)
46                x = self.conv2d(x)
47                x = self.conv_transpose2d(x)
48                return x
49
50        mm = MyModel()
51        mm.eval()
52
53        test_input = torch.ones(size=(4, 7, 5, 6), dtype=torch.float)
54
55        edge_program = capture_graph_for_xnnpack(mm, (test_input,))
56        partitioner = XnnpackFloatingPointPartitioner()
57
58        with validation_disabled():
59            delegated_program = edge_program
60            delegated_program.exported_program = to_backend(
61                edge_program.exported_program, partitioner
62            )
63
64        program = delegated_program.to_executorch(
65            get_xnnpack_executorch_backend_config([SpecPropPass()]),
66        )
67
68        size_information = generate_model_size_information(
69            model=program,
70            delegate_deserializers=None,
71            flatbuffer=program.buffer,
72        )
73
74        # Number of Elements -> Other tensor data
75        exepected_tensor_data = {
76            # Conv3d Weight
77            216: {
78                "dtype": "float32",
79                "element_size": 4,
80                "shape": [2, 4, 3, 3, 3],
81                "num_bytes": 864,
82            },
83            # ConvTranspose2d Weight
84            32: {
85                "dtype": "float32",
86                "element_size": 4,
87                "shape": [2, 4, 2, 2],
88                "num_bytes": 128,
89            },
90            # ConvTranspose2d Bias
91            4: {
92                "dtype": "float32",
93                "element_size": 4,
94                "shape": [4],
95                "num_bytes": 16,
96            },
97            # Conv3d Bias
98            2: {
99                "dtype": "float32",
100                "element_size": 4,
101                "shape": [2],
102                "num_bytes": 8,
103            },
104        }
105
106        self.assertEqual(
107            len(size_information["tensor_data"]), len(exepected_tensor_data)
108        )
109
110        for tensor in size_information["tensor_data"]:
111            for k, v in exepected_tensor_data[tensor["numel"]].items():
112                self.assertEqual(tensor[k], v)
113
114        # Two delegate blobs: sigmoid and conv2d
115        self.assertEqual(len(size_information["delegate_blob_data"]), 2)
116