Skip to content

Commit 58fa495

Browse files
committed
Update on "Enable aoti for preprocess"
Land and update torch nightly pin after: pytorch/pytorch#137063 Test Plan: With pytorch/pytorch#137063: ``` python -m unittest examples.models.llama3_2_vision.preprocess.test_preprocess ``` [ghstack-poisoned]
1 parent 4755a20 commit 58fa495

File tree

2 files changed

+17
-16
lines changed

2 files changed

+17
-16
lines changed

examples/models/llama3_2_vision/preprocess/export_preprocess.py

Lines changed: 9 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -24,29 +24,22 @@ def main():
2424
strict=False,
2525
)
2626

27-
# Executorch
27+
# AOTInductor. Note: export AOTI before ExecuTorch, as
28+
# ExecuTorch will modify the ExportedProgram.
29+
torch._inductor.aot_compile(
30+
ep.module(),
31+
model.get_example_inputs(),
32+
options={"aot_inductor.output_path": "preprocess_aoti.so"},
33+
)
34+
35+
# Executorch.
2836
edge_program = to_edge(
2937
ep, compile_config=EdgeCompileConfig(_check_ir_validity=False)
3038
)
3139
et_program = edge_program.to_executorch()
3240
with open("preprocess_et.pte", "wb") as file:
3341
et_program.write_to_file(file)
3442

35-
# Export.
36-
# ep = torch.export.export(
37-
# model.get_eager_model(),
38-
# model.get_example_inputs(),
39-
# dynamic_shapes=model.get_dynamic_shapes(),
40-
# strict=False,
41-
# )
42-
#
43-
# # AOTInductor
44-
# torch._inductor.aot_compile(
45-
# ep.module(),
46-
# model.get_example_inputs(),
47-
# options={"aot_inductor.output_path": "preprocess_aoti.so"},
48-
# )
49-
5043

5144
if __name__ == "__main__":
5245
main()

examples/models/llama3_2_vision/preprocess/test_preprocess.py

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -86,6 +86,14 @@ def initialize_models(resize_to_max_canvas: bool) -> Dict[str, Any]:
8686
)
8787
executorch_model = edge_program.to_executorch()
8888

89+
# Re-export, as lowering to executorch changes the graph.
90+
exported_model = torch.export.export(
91+
model.get_eager_model(),
92+
model.get_example_inputs(),
93+
dynamic_shapes=model.get_dynamic_shapes(),
94+
strict=False,
95+
)
96+
8997
return {
9098
"config": config,
9199
"reference_model": reference_model,

0 commit comments

Comments
 (0)