Skip to content

Commit de84b29

Browse files
guangy10facebook-github-bot
authored andcommitted
Fix export config in examples (#34)
Summary: Pull Request resolved: #34 ## Context It doesn't look correct that want to export models with one config but test the export with a different one. ## This DIff - Ensure canonical config is used for both export and tests - Ensure the test is loading the lowered model through `_load_for_executorch_from_buffer` for output consistency checking - Remove confusion about certain export flags by adding more inline comments to those flags, e.g. why is `enable_dynamic_shape` there not enabled by default, what does `enable_aot` do with dynamic shapes, etc. Reviewed By: JacobSzwejbka Differential Revision: D48018569 fbshipit-source-id: 0e56734d59b74ccf7a1fe784af78b119cdd01947
1 parent 7cffa86 commit de84b29

File tree

4 files changed

+12
-7
lines changed

4 files changed

+12
-7
lines changed

examples/export/test/TARGETS

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -10,5 +10,6 @@ python_unittest(
1010
"//executorch/examples/export:utils",
1111
"//executorch/examples/models:models",
1212
"//executorch/exir:lib",
13+
"//executorch/extension/pybindings:portable", # @manual
1314
],
1415
)

examples/export/test/test_export.py

Lines changed: 8 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -8,8 +8,11 @@
88

99
import torch
1010

11-
from executorch.examples.export.utils import _EDGE_COMPILE_CONFIG
11+
from executorch.examples.export.utils import _CAPTURE_CONFIG, _EDGE_COMPILE_CONFIG
1212
from executorch.examples.models import MODEL_NAME_TO_MODEL
13+
from executorch.extension.pybindings.portable import ( # @manual
14+
_load_for_executorch_from_buffer,
15+
)
1316

1417

1518
class ExportTest(unittest.TestCase):
@@ -18,16 +21,17 @@ def _assert_eager_lowered_same_result(
1821
):
1922
import executorch.exir as exir
2023

21-
capture_config = exir.CaptureConfig(enable_dynamic_shape=False)
22-
edge_model = exir.capture(eager_model, example_inputs, capture_config).to_edge(
24+
edge_model = exir.capture(eager_model, example_inputs, _CAPTURE_CONFIG).to_edge(
2325
_EDGE_COMPILE_CONFIG
2426
)
2527

2628
executorch_model = edge_model.to_executorch()
29+
pte_model = _load_for_executorch_from_buffer(executorch_model.buffer)
30+
2731
with torch.no_grad():
2832
eager_output = eager_model(*example_inputs)
2933
with torch.no_grad():
30-
executorch_output = executorch_model.graph_module(*example_inputs)
34+
executorch_output = pte_model.forward(example_inputs)
3135
self.assertTrue(
3236
torch.allclose(eager_output, executorch_output[0], rtol=1e-5, atol=1e-5)
3337
)

examples/export/utils.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -11,7 +11,7 @@
1111
# Reason is that there memory allocation ops with symbolic shape nodes.
1212
# and when evaulating shape, it doesnt seem that we presenting them with shape env
1313
# that contain those variables.
14-
_CAPTURE_CONFIG = exir.CaptureConfig(enable_aot=True, _unlift=False)
14+
_CAPTURE_CONFIG = exir.CaptureConfig(enable_aot=True)
1515
_EDGE_COMPILE_CONFIG = exir.EdgeCompileConfig(
1616
_check_ir_validity=False,
1717
)

exir/capture/_config.py

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -19,8 +19,8 @@
1919
class CaptureConfig:
2020
pt2_mode: bool = True
2121
enable_functionalization: bool = True
22-
enable_dynamic_shape: bool = False
23-
enable_aot: bool = False
22+
enable_dynamic_shape: bool = False # This flag does nothing if enable_aot is True
23+
enable_aot: bool = False # When it's true it implies automatic dynamic shapes via default dynamo config
2424
_dynamo_config: "ExirDynamoConfig" = field(default_factory=ExirDynamoConfig)
2525
_unlift: bool = False
2626
_use_old_decomp_table: bool = False

0 commit comments

Comments
 (0)