From 6d57cde2f1fe5ae176715d90c5420542097dcd0c Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Mon, 13 May 2024 15:20:46 -0700 Subject: [PATCH] Ease Core ML partitioner and quantizer imports. (#3564) Summary: . Reviewed By: mcr229 Differential Revision: D57172491 --- backends/apple/coreml/README.md | 4 ++-- backends/apple/coreml/partition/__init__.py | 9 +++++++++ backends/apple/coreml/quantizer/__init__.py | 9 +++++++++ backends/apple/coreml/test/test_coreml_partitioner.py | 4 +--- backends/apple/coreml/test/test_coreml_quantizer.py | 2 +- examples/apple/coreml/scripts/export.py | 4 +--- examples/models/llama2/lib/partitioner_lib.py | 8 +++----- 7 files changed, 26 insertions(+), 14 deletions(-) create mode 100644 backends/apple/coreml/partition/__init__.py create mode 100644 backends/apple/coreml/quantizer/__init__.py diff --git a/backends/apple/coreml/README.md b/backends/apple/coreml/README.md index 4a21d8d8ae1..05b56e9c788 100644 --- a/backends/apple/coreml/README.md +++ b/backends/apple/coreml/README.md @@ -28,7 +28,7 @@ import torch import executorch.exir from executorch.backends.apple.coreml.compiler import CoreMLBackend -from executorch.backends.apple.coreml.partition.coreml_partitioner import CoreMLPartitioner +from executorch.backends.apple.coreml.partition import CoreMLPartitioner class Model(torch.nn.Module): def __init__(self): @@ -72,7 +72,7 @@ from torch.ao.quantization.quantize_pt2e import ( prepare_qat_pt2e, ) -from executorch.backends.apple.coreml.quantizer.coreml_quantizer import CoreMLQuantizer +from executorch.backends.apple.coreml.quantizer import CoreMLQuantizer from coremltools.optimize.torch.quantization.quantization_config import ( LinearQuantizerConfig, QuantizationScheme, diff --git a/backends/apple/coreml/partition/__init__.py b/backends/apple/coreml/partition/__init__.py new file mode 100644 index 00000000000..1630e9ece45 --- /dev/null +++ b/backends/apple/coreml/partition/__init__.py @@ -0,0 +1,9 @@ +# Copyright © 2023 Apple Inc. All rights reserved. +# +# Please refer to the license found in the LICENSE file in the root directory of the source tree. + +from .coreml_partitioner import CoreMLPartitioner + +__all__ = [ + CoreMLPartitioner, +] diff --git a/backends/apple/coreml/quantizer/__init__.py b/backends/apple/coreml/quantizer/__init__.py new file mode 100644 index 00000000000..f6282834fa1 --- /dev/null +++ b/backends/apple/coreml/quantizer/__init__.py @@ -0,0 +1,9 @@ +# Copyright © 2023 Apple Inc. All rights reserved. +# +# Please refer to the license found in the LICENSE file in the root directory of the source tree. + +from .coreml_quantizer import CoreMLQuantizer + +__all__ = [ + CoreMLQuantizer, +] diff --git a/backends/apple/coreml/test/test_coreml_partitioner.py b/backends/apple/coreml/test/test_coreml_partitioner.py index e59e5c95544..45c468e450b 100644 --- a/backends/apple/coreml/test/test_coreml_partitioner.py +++ b/backends/apple/coreml/test/test_coreml_partitioner.py @@ -9,9 +9,7 @@ import torch import torchvision -from executorch.backends.apple.coreml.partition.coreml_partitioner import ( - CoreMLPartitioner, -) +from executorch.backends.apple.coreml.partition import CoreMLPartitioner class TestCoreMLPartitioner(unittest.TestCase): diff --git a/backends/apple/coreml/test/test_coreml_quantizer.py b/backends/apple/coreml/test/test_coreml_quantizer.py index 67eee3593fd..c05cde05a0a 100644 --- a/backends/apple/coreml/test/test_coreml_quantizer.py +++ b/backends/apple/coreml/test/test_coreml_quantizer.py @@ -14,7 +14,7 @@ QuantizationScheme, ) -from executorch.backends.apple.coreml.quantizer.coreml_quantizer import CoreMLQuantizer +from executorch.backends.apple.coreml.quantizer import CoreMLQuantizer from torch._export import capture_pre_autograd_graph from torch.ao.quantization.quantize_pt2e import ( convert_pt2e, diff --git a/examples/apple/coreml/scripts/export.py b/examples/apple/coreml/scripts/export.py index 966714ba31c..4bf26a7f3ea 100644 --- a/examples/apple/coreml/scripts/export.py +++ b/examples/apple/coreml/scripts/export.py @@ -16,9 +16,7 @@ from executorch.backends.apple.coreml.compiler import CoreMLBackend -from executorch.backends.apple.coreml.partition.coreml_partitioner import ( - CoreMLPartitioner, -) +from executorch.backends.apple.coreml.partition import CoreMLPartitioner from executorch.exir import to_edge from executorch.exir.backend.backend_api import to_backend diff --git a/examples/models/llama2/lib/partitioner_lib.py b/examples/models/llama2/lib/partitioner_lib.py index 1638a357576..c11e74a7e0b 100644 --- a/examples/models/llama2/lib/partitioner_lib.py +++ b/examples/models/llama2/lib/partitioner_lib.py @@ -57,16 +57,14 @@ def get_coreml_partitioner(args): args.use_kv_cache is True ), "CoreML backend currently only supports static shape and use_kv_cache=True is the only way to support it at the moment" try: - # pyre-ignore: Undefined import [21]: Could not find a module corresponding to import `executorch.backends.apple.coreml.partition.coreml_partitioner`. + # pyre-ignore: Undefined import [21]: Could not find a module corresponding to import `coremltools`. import coremltools as ct # pyre-ignore: Undefined import [21]: Could not find a module corresponding to import `executorch.backends.apple.coreml.compiler` from executorch.backends.apple.coreml.compiler import CoreMLBackend - # pyre-ignore: Undefined import [21]: Could not find a module corresponding to import `executorch.backends.apple.coreml.partition.coreml_partitioner` - from executorch.backends.apple.coreml.partition.coreml_partitioner import ( - CoreMLPartitioner, - ) + # pyre-ignore: Undefined import [21]: Could not find a module corresponding to import `executorch.backends.apple.coreml.partition` + from executorch.backends.apple.coreml.partition import CoreMLPartitioner except ImportError: raise ImportError( "Please install the CoreML backend follwing https://pytorch.org/executorch/main/build-run-coreml.html"