Coreml recipe updates (#108)

* init

* up

* up

* up

Author: metascroy

optimum/exporters/executorch/recipes/coreml.py

@@ -13,7 +13,8 @@
 # limitations under the License.
 
 import logging
-from typing import Dict, Union
+from itertools import product
+from typing import Any, Dict, Union
 
 from tabulate import tabulate
 from torch.export import ExportedProgram
@@ -34,8 +35,7 @@
 from ..recipe_registry import register_recipe
 
 
-@register_recipe("coreml")
-def export_to_executorch_with_coreml(
+def _export_to_executorch(
     model: Union[CausalLMExportableModule, MaskedLMExportableModule, Seq2SeqLMExportableModule],
     **kwargs,
 ):
@@ -63,23 +63,14 @@ def export_to_executorch_with_coreml(
 
     def _lower_to_executorch(
         exported_programs: Dict[str, ExportedProgram],
-        metadata=None,
-        **kwargs,
+        metadata,
+        compute_unit,
+        minimum_deployment_target,
+        compute_precision,
     ) -> Dict[str, ExecutorchProgram]:
-        compute_unit = kwargs.get("compute_unit", ct.ComputeUnit.ALL)
-        minimum_deployment_target = kwargs.get("minimum_deployment_target", ct.target.iOS15)
-        compute_precision = kwargs.get("compute_precision", ct.precision.FLOAT16)
-        model_type = kwargs.get("model_type", "model")
-        model_type = {
-            "model": CoreMLBackend.MODEL_TYPE.MODEL,
-            "modelc": CoreMLBackend.MODEL_TYPE.COMPILED_MODEL,
-        }[model_type]
-        take_over_mutable_buffer = kwargs.get("take_over_mutable_buffer", True)
-
         et_progs = {}
         backend_config_dict = {}
         for pte_name, exported_program in exported_programs.items():
-            exported_program = exported_program.run_decompositions({})
             logging.debug(f"\nExported program for {pte_name}.pte: {exported_program}")
             et_progs[pte_name] = to_edge_transform_and_lower(
                 exported_program,
@@ -89,14 +80,15 @@ def _lower_to_executorch(
                             compute_unit=compute_unit,
                             minimum_deployment_target=minimum_deployment_target,
                             compute_precision=compute_precision,
-                            model_type=model_type,
+                            model_type=CoreMLBackend.MODEL_TYPE.MODEL,
                         ),
-                        take_over_mutable_buffer=take_over_mutable_buffer,
+                        take_over_mutable_buffer=(minimum_deployment_target >= ct.target.iOS18),
                     )
                 ],
                 compile_config=EdgeCompileConfig(
                     _check_ir_validity=False,
-                    _skip_dim_order=False,
+                    # In ET 0.7, we can set _skip_dim_order=False
+                    _skip_dim_order=True,
                 ),
                 constant_methods=metadata,
             ).to_executorch(
@@ -114,3 +106,46 @@ def _lower_to_executorch(
 
     exported_progs = model.export()
     return _lower_to_executorch(exported_progs, model.metadata, **kwargs)
+
+
+def _get_recipe_kwargs(dtype: str, compute_unit: str) -> Dict[str, Any]:
+    import coremltools as ct
+
+    compute_precision = {
+        "fp16": ct.precision.FLOAT16,
+        "fp32": ct.precision.FLOAT32,
+    }[dtype]
+
+    compute_unit = {
+        "cpu": ct.ComputeUnit.CPU_ONLY,
+        "gpu": ct.ComputeUnit.CPU_AND_GPU,
+        "ne": ct.ComputeUnit.CPU_AND_NE,
+        "all": ct.ComputeUnit.ALL,
+    }[compute_unit]
+
+    recipe_kwargs = {
+        "compute_precision": compute_precision,
+        "compute_unit": compute_unit,
+        "minimum_deployment_target": ct.target.iOS18,
+    }
+    return recipe_kwargs
+
+
+def _make_recipe(recipe_name, recipe_kwargs):
+    @register_recipe(recipe_name)
+    def recipe_fn(exported_programs: Dict[str, ExportedProgram], **kwargs):
+        return _export_to_executorch(
+            exported_programs,
+            **recipe_kwargs,
+        )
+
+    return recipe_fn
+
+
+# Register recipes for CoreML backend
+for dtype, compute_unit in product(["fp32", "fp16"], ["cpu", "gpu", "ne", "all"]):
+    recipe_name = f"coreml_{dtype}"
+    if compute_unit != "all":
+        recipe_name += f"_{compute_unit}"
+    recipe_kwargs = _get_recipe_kwargs(dtype=dtype, compute_unit=compute_unit)
+    _make_recipe(recipe_name, recipe_kwargs)

tests/models/test_modeling_vit.py

@@ -89,30 +89,4 @@ def test_vit_image_classification_portable(self):
     @pytest.mark.run_slow
     @pytest.mark.skipif(is_not_macos, reason="Only runs on MacOS")
     def test_vit_image_classification_coreml_fp32_cpu(self):
-        model_id = "google/vit-base-patch16-224"
-
-        config = AutoConfig.from_pretrained(model_id)
-        batch_size = 1
-        num_channels = config.num_channels
-        height = config.image_size
-        width = config.image_size
-        pixel_values = torch.rand(batch_size, num_channels, height, width)
-
-        # Test fetching and lowering the model to ExecuTorch
-        import coremltools as ct
-
-        et_model = ExecuTorchModelForImageClassification.from_pretrained(
-            model_id=model_id,
-            recipe="coreml",
-            recipe_kwargs={"compute_precision": ct.precision.FLOAT32, "compute_units": ct.ComputeUnit.CPU_ONLY},
-        )
-        self.assertIsInstance(et_model, ExecuTorchModelForImageClassification)
-        self.assertIsInstance(et_model.model, ExecuTorchModule)
-
-        eager_model = AutoModelForImageClassification.from_pretrained(model_id).eval().to("cpu")
-        with torch.no_grad():
-            eager_output = eager_model(pixel_values)
-            et_output = et_model.forward(pixel_values)
-
-        # Compare with eager outputs
-        self.assertTrue(check_close_recursively(eager_output.logits, et_output))
+        self._helper_vit_image_classification(recipe="coreml_fp32_cpu")