Apple Silicon Support (ORT and CoreML)

README.md

@@ -2,7 +2,7 @@
 
 [![Turnkey tests](https://github.com/onnx/turnkeyml/actions/workflows/test_turnkey.yml/badge.svg)](https://github.com/onnx/turnkeyml/tree/main/test "Check out our tests")
 [![Build API tests](https://github.com/onnx/turnkeyml/actions/workflows/test_build_api.yml/badge.svg)](https://github.com/onnx/turnkeyml/tree/main/test "Check out our tests")
-[![OS - Windows | Linux](https://img.shields.io/badge/OS-windows%20%7C%20linux-blue)](https://github.com/onnx/turnkeyml/blob/main/docs/install.md "Check out our instructions")
+[![OS - Windows | Linux | MacOS](https://img.shields.io/badge/OS-Windows%20%7C%20Linux%20%7C%20MacOS-blue)](https://github.com/onnx/turnkeyml/blob/main/docs/install.md "Check out our instructions")
 [![Made with Python](https://img.shields.io/badge/Python-3.8,3.10-blue?logo=python&logoColor=white)](https://github.com/onnx/turnkeyml/blob/main/docs/install.md "Check out our instructions")
 
 

docs/install.md

@@ -20,6 +20,13 @@ bash Miniconda3-latest-Linux-x86_64.sh
 
 If you are installing TurnkeyML on **Windows**, manually download and install [Miniconda3 for Windows 64-bit](https://repo.anaconda.com/miniconda/Miniconda3-latest-Windows-x86_64.exe). Please note that PowerShell is recommended when using miniconda on Windows.
 
+
+If you are installing TurnkeyML on **MacOS**, run the command below:
+```
+wget https://repo.anaconda.com/miniconda/Miniconda3-latest-MacOSX-x86_64.sh
+Miniconda3-latest-MacOSX-x86_64.sh
+```
+
 Then create and activate a virtual environment like this:
 
 ```

docs/tools_user_guide.md

@@ -12,11 +12,14 @@ The tools currently support the following combinations of runtimes and devices:
 | ----------- | ---------- | ------------------------------------------------------------------------------------- | -------------------------------- | --------------------------------------------- |
 | Nvidia GPU  | nvidia     | TensorRT<sup>†</sup>                                                                  | trt                              | Any Nvidia GPU supported by TensorRT          |
 | x86 CPU     | x86        | ONNX Runtime<sup>‡</sup>, Pytorch Eager, Pytoch 2.x Compiled | ort, torch-eager, torch-compiled | Any Intel or AMD CPU supported by the runtime |
+| Apple Silicon    | apple_silicon        | CoreML<sup>◊</sup>, ONNX Runtime<sup>‡</sup>, Pytorch Eager | coreml, ort, torch-eager | Any Apple M* Silicon supported by the runtime |
+
 </span>
 
 <sup>†</sup> Requires TensorRT >= 8.5.2  
 <sup>‡</sup> Requires ONNX Runtime >= 1.13.1  
 <sup>*</sup> Requires Pytorch >= 2.0.0  
+<sup>◊</sup> Requires CoreML >= 7.1  
 
 
 # Table of Contents
@@ -252,6 +255,9 @@ Each device type has its own default runtime, as indicated below.
   - `torch-compiled`: PyTorch 2.x-style compiled graph execution using TorchInductor.
 - Valid runtimes for `nvidia` device
   - `trt`: Nvidia TensorRT (default).
+- Valid runtimes for `apple_silicon` device
+  - `coreml`: CoreML (default).
+  - `ort`: ONNX Runtime.
 
 This feature is also be available as an API argument: 
 - `benchmark_files(runtime=[...])`

setup.py

@@ -18,6 +18,7 @@
         "turnkeyml.run.onnxrt",
         "turnkeyml.run.tensorrt",
         "turnkeyml.run.torchrt",
+        "turnkeyml.run.coreml",
         "turnkeyml.cli",
         "turnkeyml.common",
         "turnkeyml_models",
@@ -46,6 +47,7 @@
         "pandas>=1.5.3",
         "fasteners",
         "GitPython>=3.1.40",
+        "coremltools>=7.1",
     ],
     extras_require={
         "tensorflow": [

src/turnkeyml/build/export.py

@@ -10,6 +10,7 @@
 import onnxruntime
 import onnxmltools
 import onnx
+import coremltools as ct
 import turnkeyml.build.stage as stage
 import turnkeyml.common.exceptions as exp
 import turnkeyml.common.build as build
@@ -40,6 +41,20 @@ def _warn_to_stdout(message, category, filename, line_number, _, line):
     )
 
 
+def validate_torch_args(state: build.State) -> None:
+    """
+    Ensure that the inputs received match the model's forward function
+    """
+    all_args = list(inspect.signature(state.model.forward).parameters.keys())
+    for inp in list(state.inputs.keys()):
+        if inp not in all_args:
+            msg = f"""
+            Input name {inp} not found in the model's forward method. Available
+            input names are: {all_args}"
+            """
+            raise ValueError(msg)
+
+
 def get_output_names(
     onnx_model: Union[str, onnx.ModelProto]
 ):  # pylint: disable=no-member
@@ -62,6 +77,13 @@ def base_onnx_file(state: build.State):
     )
 
 
+def base_coreml_file(state: build.State):
+    return os.path.join(
+        onnx_dir(state),
+        f"{state.config.build_name}-op{state.config.onnx_opset}-base.mlmodel",
+    )
+
+
 def opt_onnx_file(state: build.State):
     return os.path.join(
         onnx_dir(state),
@@ -234,16 +256,7 @@ def fire(self, state: build.State):
         user_provided_args = list(state.inputs.keys())
 
         if isinstance(state.model, torch.nn.Module):
-            # Validate user provided args
-            all_args = list(inspect.signature(state.model.forward).parameters.keys())
-
-            for inp in user_provided_args:
-                if inp not in all_args:
-                    msg = f"""
-                    Input name {inp} not found in the model's forward method. Available
-                    input names are: {all_args}"
-                    """
-                    raise ValueError(msg)
+            validate_torch_args(state)
 
             # Most pytorch models have args that are kind = positional_or_keyword.
             # The `torch.onnx.export()` function accepts model args as
@@ -252,6 +265,7 @@ def fire(self, state: build.State):
             # the order of the input_names must reflect the order of the model args.
 
             # Collect order of pytorch model args.
+            all_args = list(inspect.signature(state.model.forward).parameters.keys())
             all_args_order_mapping = {arg: idx for idx, arg in enumerate(all_args)}
 
             # Sort the user provided inputs with respect to model args and store as tuple.
@@ -620,3 +634,69 @@ def fire(self, state: build.State):
             raise exp.StageError(msg)
 
         return state
+
+
+class ExportToCoreML(stage.Stage):
+    """
+    Stage that takes a Pytorch model and inputs and converts to CoreML format.
+
+    Expected inputs:
+     - state.model is a torch.nn.Module or torch.jit.ScriptModule
+     - state.inputs is a dict that represents valid kwargs to the forward
+        function of state.model
+    Outputs:
+     - A *.mlmodel file
+    """
+
+    def __init__(self):
+        super().__init__(
+            unique_name="coreml_conversion",
+            monitor_message="Converting to CoreML",
+        )
+
+    def fire(self, state: build.State):
+        if not isinstance(state.model, (torch.nn.Module, torch.jit.ScriptModule)):
+            msg = f"""
+            The current stage (ExportToCoreML) is only compatible with
+            models of type torch.nn.Module or torch.jit.ScriptModule, however
+            the stage received a model of type {type(state.model)}.
+            """
+            raise exp.StageError(msg)
+
+        if isinstance(state.model, torch.nn.Module):
+            validate_torch_args(state)
+
+        # Send warnings to stdout (and therefore the log file)
+        default_warnings = warnings.showwarning
+        warnings.showwarning = _warn_to_stdout
+
+        # Generate a TorchScript Version
+        dummy_inputs = copy.deepcopy(state.inputs)
+        traced_model = torch.jit.trace(state.model, example_kwarg_inputs=dummy_inputs)
+
+        # Export the model to CoreML
+        output_path = base_coreml_file(state)
+        os.makedirs(onnx_dir(state), exist_ok=True)
+        coreml_model = ct.convert(
+            traced_model,
+            inputs=[ct.TensorType(shape=inp.shape) for inp in dummy_inputs.values()],
+            convert_to="neuralnetwork",
+        )
+
+        # Save the CoreML model
+        coreml_model.save(output_path)
+
+        # Save output names to ensure we are preserving the order of the outputs
+        state.expected_output_names = get_output_names(output_path)
+
+        # Restore default warnings behavior
+        warnings.showwarning = default_warnings
+
+        tensor_helpers.save_inputs(
+            [state.inputs], state.original_inputs_file, downcast=False
+        )
+
+        # Save intermediate results
+        state.intermediate_results = [output_path]
+
+        return state

src/turnkeyml/build/sequences.py

@@ -32,6 +32,15 @@
     enable_model_validation=True,
 )
 
+coreml = stage.Sequence(
+    "coreml",
+    "CoreML Sequence",
+    [
+        export.ExportToCoreML(),
+    ],
+    enable_model_validation=True,
+)
+
 # Plugin interface for sequences
 discovered_plugins = plugins.discover()
 
@@ -40,6 +49,7 @@
     "optimize-fp16": optimize_fp16,
     "optimize-fp32": optimize_fp32,
     "onnx-fp32": onnx_fp32,
+    "coreml": coreml,
 }
 
 # Add sequences from plugins to supported sequences dict

src/turnkeyml/run/basert.py

@@ -53,6 +53,8 @@ def __init__(
         requires_docker: bool = False,
         tensor_type=np.array,
         execute_function: Optional[callable] = None,
+        model_filename="model.onnx",
+        model_dirname="onnxmodel",
     ):
         self.tensor_type = tensor_type
         self.cache_dir = cache_dir
@@ -71,16 +73,16 @@ def __init__(
         self.runtime_version = runtime_version
         self.model = model
         self.inputs = inputs
-        self.onnx_filename = "model.onnx"
-        self.onnx_dirname = "onnxmodel"
+        self.model_filename = model_filename
+        self.model_dirname = model_dirname
         self.outputs_filename = "outputs.json"
         self.runtimes_supported = runtimes_supported
         self.execute_function = execute_function
 
         # Validate runtime is supported
         if runtime not in runtimes_supported:
             raise ValueError(
-                f"'runtime' argument {runtime} passed to TensorRT, which only "
+                f"'runtime' argument {runtime} passed to a runtime that only "
                 f"supports runtimes: {runtimes_supported}"
             )
 
@@ -102,23 +104,23 @@ def local_output_dir(self):
         )
 
     @property
-    def local_onnx_dir(self):
-        return os.path.join(self.local_output_dir, self.onnx_dirname)
+    def local_model_dir(self):
+        return os.path.join(self.local_output_dir, self.model_dirname)
 
     @property
     def docker_onnx_dir(self):
         return self.posix_path_format(
-            os.path.join(self.docker_output_dir, self.onnx_dirname)
+            os.path.join(self.docker_output_dir, self.model_dirname)
         )
 
     @property
-    def local_onnx_file(self):
-        return os.path.join(self.local_onnx_dir, self.onnx_filename)
+    def local_model_file(self):
+        return os.path.join(self.local_model_dir, self.model_filename)
 
     @property
-    def docker_onnx_file(self):
+    def docker_model_file(self):
         return self.posix_path_format(
-            os.path.join(self.docker_onnx_dir, self.onnx_filename)
+            os.path.join(self.docker_onnx_dir, self.model_filename)
         )
 
     @property
@@ -183,16 +185,16 @@ def benchmark(self) -> MeasuredPerformance:
             raise exp.ModelRuntimeError(msg)
 
         os.makedirs(self.local_output_dir, exist_ok=True)
-        os.makedirs(self.local_onnx_dir, exist_ok=True)
-        shutil.copy(model_file, self.local_onnx_file)
+        os.makedirs(self.local_model_dir, exist_ok=True)
+        shutil.copy(model_file, self.local_model_file)
 
         # Execute benchmarking in hardware
         if self.requires_docker:
             _check_docker_install()
-            onnx_file = self.docker_onnx_file
+            onnx_file = self.docker_model_file
             _check_docker_running()
         else:
-            onnx_file = self.local_onnx_file
+            onnx_file = self.local_model_file
 
         self._execute(
             output_dir=self.local_output_dir,

src/turnkeyml/run/coreml/__init__.py

@@ -0,0 +1,13 @@
+import turnkeyml.build.sequences as sequences
+from .runtime import CoreML
+
+implements = {
+    "runtimes": {
+        "coreml": {
+            "build_required": True,
+            "RuntimeClass": CoreML,
+            "supported_devices": {"apple_silicon"},
+            "default_sequence": sequences.coreml,
+        }
+    }
+}