Uh oh!
There was an error while loading. Please reload this page.
forked from pytorch/executorch
- Notifications
You must be signed in to change notification settings - Fork 0
Expand file tree
/
Copy pathexport_bundled_program.py
More file actions
Latest commit
150 lines (119 loc) · 4.86 KB
/
Copy pathexport_bundled_program.py
File metadata and controls
150 lines (119 loc) · 4.86 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
# 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.
# Example script for exporting simple models to flatbuffer
# pyre-unsafe
importargparse
fromtypingimportList
importtorch
fromexecutorch.devtoolsimportBundledProgram
fromexecutorch.devtools.bundled_program.configimport (
MethodInputType,
MethodTestCase,
MethodTestSuite,
)
fromexecutorch.devtools.bundled_program.serializeimport (
serialize_from_bundled_program_to_flatbuffer,
)
fromexecutorch.exirimportExecutorchProgramManager
fromexecutorch.extension.export_util.utilsimportexport_to_exec_prog
from ...modelsimportMODEL_NAME_TO_MODEL
from ...models.model_factoryimportEagerModelFactory
defsave_bundled_program(
executorch_program: ExecutorchProgramManager,
method_test_suites: List[MethodTestSuite],
output_path: str,
):
"""
Generates a bundled program from the given ET program and saves it to the specified path.
Args:
executorch_program: The ExecuTorch program to bundle.
method_test_suites: The MethodTestSuites which contains test cases to include in the bundled program.
output_path: Path to save the bundled program.
"""
bundled_program=BundledProgram(executorch_program, method_test_suites)
bundled_program_buffer=serialize_from_bundled_program_to_flatbuffer(
bundled_program
)
withopen(output_path, "wb") asfile:
file.write(bundled_program_buffer)
defexport_to_bundled_program(
model_name: str,
output_directory: str,
model: torch.nn.Module,
example_inputs: MethodInputType,
) ->None:
"""
Exports the given eager model to bundled program.
Args:
model_name: Name of the bundled program to export.
output_directory: Directory where the bundled program should be saved.
model: The eager model to export.
example_inputs: An example input for model's all method for single execution.
To simplify, here we assume that all inference methods have the same inputs.
"""
print("Exporting ET program...")
# pyre-ignore[6]
executorch_program=export_to_exec_prog(model, example_inputs)
print("Creating bundled test cases...")
method_names= [
method.nameformethodinexecutorch_program.executorch_program.execution_plan
]
# A model could have multiple entry point methods and each of them can have multiple inputs bundled for testing.
# This example demonstrates a model which has multiple entry point methods, whose name listed in method_names, to which we want
# to bundle two input test cases (example_inputs is used two times) for each inference method.
program_inputs= {
m_name: [example_inputs, example_inputs] form_nameinmethod_names
}
method_test_suites: List[MethodTestSuite] = []
form_nameinmethod_names:
method_inputs=program_inputs[m_name]
# To create a bundled program, we first create every test cases from input. We leverage eager model
# to generate expected output for each test input, and use MethodTestCase to hold the information of
# each test case. We gather all MethodTestCase for same method into one MethodTestSuite, and generate
# bundled program by all MethodTestSuites.
method_test_cases: List[MethodTestCase] = []
formethod_inputinmethod_inputs:
method_test_cases.append(
MethodTestCase(
inputs=method_input,
expected_outputs=model(*method_input),
)
)
method_test_suites.append(
MethodTestSuite(
method_name=m_name,
test_cases=method_test_cases,
)
)
save_bundled_program(
executorch_program, method_test_suites, f"{model_name}_bundled.bpte"
)
defmain() ->None:
parser=argparse.ArgumentParser()
parser.add_argument(
"-m",
"--model_name",
required=True,
help=f"provide a model name. Valid ones: {list(MODEL_NAME_TO_MODEL.keys())}",
)
parser.add_argument(
"-d",
"--dir",
default=".",
help="the directory to store the exported bundled program. Default is current directory.",
)
args=parser.parse_args()
ifargs.model_namenotinMODEL_NAME_TO_MODEL:
raiseRuntimeError(
f"Model {args.model_name} is not a valid name. "
f"Available models are {list(MODEL_NAME_TO_MODEL.keys())}."
)
model, example_inputs, _, _=EagerModelFactory.create_model(
*MODEL_NAME_TO_MODEL[args.model_name]
)
export_to_bundled_program(args.model_name, args.dir, model, example_inputs)
if__name__=="__main__":
main() # pragma: no cover