xref: /aosp_15_r20/external/executorch/examples/devtools/scripts/export_bundled_program.py (revision 523fa7a60841cd1ecfb9cc4201f1ca8b03ed023a) 
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
7# Example script for exporting simple models to flatbuffer
8
9# pyre-unsafe
10
11import argparse
12
13from typing import List
14
15import torch
16from executorch.devtools import BundledProgram
17from executorch.devtools.bundled_program.config import (
18    MethodInputType,
19    MethodTestCase,
20    MethodTestSuite,
21)
22from executorch.devtools.bundled_program.serialize import (
23    serialize_from_bundled_program_to_flatbuffer,
24)
25
26from executorch.exir import ExecutorchProgramManager
27from executorch.extension.export_util.utils import export_to_exec_prog
28
29from ...models import MODEL_NAME_TO_MODEL
30from ...models.model_factory import EagerModelFactory
31
32
33def save_bundled_program(
34    executorch_program: ExecutorchProgramManager,
35    method_test_suites: List[MethodTestSuite],
36    output_path: str,
37):
38    """
39    Generates a bundled program from the given ET program and saves it to the specified path.
40
41    Args:
42        executorch_program: The ExecuTorch program to bundle.
43        method_test_suites: The MethodTestSuites which contains test cases to include in the bundled program.
44        output_path: Path to save the bundled program.
45    """
46
47    bundled_program = BundledProgram(executorch_program, method_test_suites)
48    bundled_program_buffer = serialize_from_bundled_program_to_flatbuffer(
49        bundled_program
50    )
51
52    with open(output_path, "wb") as file:
53        file.write(bundled_program_buffer)
54
55
56def export_to_bundled_program(
57    model_name: str,
58    output_directory: str,
59    model: torch.nn.Module,
60    example_inputs: MethodInputType,
61) -> None:
62    """
63    Exports the given eager model to bundled program.
64
65    Args:
66        model_name: Name of the bundled program to export.
67        output_directory: Directory where the bundled program should be saved.
68        model: The eager model to export.
69        example_inputs: An example input for model's all method for single execution.
70                        To simplify, here we assume that all inference methods have the same inputs.
71    """
72
73    print("Exporting ET program...")
74
75    # pyre-ignore[6]
76    executorch_program = export_to_exec_prog(model, example_inputs)
77
78    print("Creating bundled test cases...")
79    method_names = [
80        method.name for method in executorch_program.executorch_program.execution_plan
81    ]
82
83    # A model could have multiple entry point methods and each of them can have multiple inputs bundled for testing.
84    # This example demonstrates a model which has multiple entry point methods, whose name listed in method_names, to which we want
85    # to bundle two input test cases (example_inputs is used two times) for each inference method.
86    program_inputs = {
87        m_name: [example_inputs, example_inputs] for m_name in method_names
88    }
89
90    method_test_suites: List[MethodTestSuite] = []
91    for m_name in method_names:
92        method_inputs = program_inputs[m_name]
93
94        # To create a bundled program, we first create every test cases from input. We leverage eager model
95        # to generate expected output for each test input, and use MethodTestCase to hold the information of
96        # each test case. We gather all MethodTestCase for same method into one MethodTestSuite, and generate
97        # bundled program by all MethodTestSuites.
98        method_test_cases: List[MethodTestCase] = []
99        for method_input in method_inputs:
100            method_test_cases.append(
101                MethodTestCase(
102                    inputs=method_input,
103                    expected_outputs=model(*method_input),
104                )
105            )
106
107        method_test_suites.append(
108            MethodTestSuite(
109                method_name=m_name,
110                test_cases=method_test_cases,
111            )
112        )
113
114    save_bundled_program(
115        executorch_program, method_test_suites, f"{model_name}_bundled.bpte"
116    )
117
118
119def main() -> None:
120    parser = argparse.ArgumentParser()
121    parser.add_argument(
122        "-m",
123        "--model_name",
124        required=True,
125        help=f"provide a model name. Valid ones: {list(MODEL_NAME_TO_MODEL.keys())}",
126    )
127    parser.add_argument(
128        "-d",
129        "--dir",
130        default=".",
131        help="the directory to store the exported bundled program. Default is current directory.",
132    )
133
134    args = parser.parse_args()
135
136    if args.model_name not in MODEL_NAME_TO_MODEL:
137        raise RuntimeError(
138            f"Model {args.model_name} is not a valid name. "
139            f"Available models are {list(MODEL_NAME_TO_MODEL.keys())}."
140        )
141
142    model, example_inputs, _, _ = EagerModelFactory.create_model(
143        *MODEL_NAME_TO_MODEL[args.model_name]
144    )
145
146    export_to_bundled_program(args.model_name, args.dir, model, example_inputs)
147
148
149if __name__ == "__main__":
150    main()  # pragma: no cover
151