Update

[ghstack-poisoned]

Author: GregoryComer

.ci/scripts/wheel/test_windows.py

@@ -5,84 +5,74 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
-try:
-    from typing import List
-
-    print("Dumping env: ", file=sys.stderr)
-    import os
-    for name, value in os.environ.items():
-        print(f"  {name}: {value}")
-
-    import sys
-    print("Trying to import ET...", file=sys.stderr)
-    import executorch
-    print(f"ET path: {executorch.__path__}", file=sys.stderr)
-
-    import torch
-    from executorch.backends.xnnpack.partition.xnnpack_partitioner import XnnpackPartitioner
-    from executorch.examples.models import Backend, Model, MODEL_NAME_TO_MODEL
-    from executorch.examples.models.model_factory import EagerModelFactory
-    from executorch.examples.xnnpack import MODEL_NAME_TO_OPTIONS
-    from executorch.examples.xnnpack.quantization.utils import quantize as quantize_xnn
-    from executorch.exir import EdgeCompileConfig, to_edge_transform_and_lower
-    from executorch.extension.pybindings.portable_lib import (
-        _load_for_executorch_from_buffer,
+from typing import List
+
+import sys
+import executorch
+
+import torch
+from executorch.backends.xnnpack.partition.xnnpack_partitioner import XnnpackPartitioner
+from executorch.examples.models import Backend, Model, MODEL_NAME_TO_MODEL
+from executorch.examples.models.model_factory import EagerModelFactory
+from executorch.examples.xnnpack import MODEL_NAME_TO_OPTIONS
+from executorch.examples.xnnpack.quantization.utils import quantize as quantize_xnn
+from executorch.exir import EdgeCompileConfig, to_edge_transform_and_lower
+from executorch.extension.pybindings.portable_lib import (
+    _load_for_executorch_from_buffer,
+)
+from test_base import ModelTest
+
+
+def test_model_xnnpack(model: Model, quantize: bool) -> None:
+    model_instance, example_inputs, _, _ = EagerModelFactory.create_model(
+        *MODEL_NAME_TO_MODEL[str(model)]
     )
-    from test_base import ModelTest
 
+    model_instance.eval()
+    ref_outputs = model_instance(*example_inputs)
 
-    def test_model_xnnpack(model: Model, quantize: bool) -> None:
-        model_instance, example_inputs, _, _ = EagerModelFactory.create_model(
-            *MODEL_NAME_TO_MODEL[str(model)]
+    if quantize:
+        quant_type = MODEL_NAME_TO_OPTIONS[str(model)].quantization
+        model_instance = torch.export.export_for_training(
+            model_instance, example_inputs
+        )
+        model_instance = quantize_xnn(
+            model_instance.module(), example_inputs, quant_type
         )
 
-        model_instance.eval()
-        ref_outputs = model_instance(*example_inputs)
-
-        if quantize:
-            quant_type = MODEL_NAME_TO_OPTIONS[str(model)].quantization
-            model_instance = torch.export.export_for_training(
-                model_instance, example_inputs
-            )
-            model_instance = quantize_xnn(
-                model_instance.module(), example_inputs, quant_type
-            )
-
-        lowered = to_edge_transform_and_lower(
-            torch.export.export(model_instance, example_inputs),
-            partitioner=[XnnpackPartitioner()],
-            compile_config=EdgeCompileConfig(
-                _check_ir_validity=False,
-            ),
-        ).to_executorch()
+    lowered = to_edge_transform_and_lower(
+        torch.export.export(model_instance, example_inputs),
+        partitioner=[XnnpackPartitioner()],
+        compile_config=EdgeCompileConfig(
+            _check_ir_validity=False,
+        ),
+    ).to_executorch()
 
-        loaded_model = _load_for_executorch_from_buffer(lowered.buffer)
-        et_outputs = loaded_model([*example_inputs])
+    loaded_model = _load_for_executorch_from_buffer(lowered.buffer)
+    et_outputs = loaded_model([*example_inputs])
 
-        if isinstance(ref_outputs, torch.Tensor):
-            ref_outputs = (ref_outputs,)
+    if isinstance(ref_outputs, torch.Tensor):
+        ref_outputs = (ref_outputs,)
 
-        assert len(ref_outputs) == len(et_outputs)
-        for i in range(len(ref_outputs)):
-            assert torch.allclose(ref_outputs[i], et_outputs[i], atol=1e-5)
+    assert len(ref_outputs) == len(et_outputs)
+    for i in range(len(ref_outputs)):
+        assert torch.allclose(ref_outputs[i], et_outputs[i], atol=1e-5)
 
 
-    def run_tests(model_tests: List[ModelTest]) -> None:
-        for model_test in model_tests:
-            if model_test.backend == Backend.Xnnpack:
-                test_model_xnnpack(model_test.model, quantize=False)
-            else:
-                raise RuntimeError(f"Unsupported backend {model_test.backend}.")
+def run_tests(model_tests: List[ModelTest]) -> None:
+    for model_test in model_tests:
+        if model_test.backend == Backend.Xnnpack:
+            test_model_xnnpack(model_test.model, quantize=False)
+        else:
+            raise RuntimeError(f"Unsupported backend {model_test.backend}.")
 
 
-    if __name__ == "__main__":
-        run_tests(
-            model_tests=[
-                ModelTest(
-                    model=Model.Mv3,
-                    backend=Backend.Xnnpack,
-                ),
-            ]
-        )
-except:
-    pass
+if __name__ == "__main__":
+    run_tests(
+        model_tests=[
+            ModelTest(
+                model=Model.Mv3,
+                backend=Backend.Xnnpack,
+            ),
+        ]
+    )