Demo CoreML Partitioner (#1981)
Summary:
## TLDR
This PR adds ViT as an e2e example of leveraging CoreML partitioner
💁♂️ Tips for reviewers: Review commit by commit would be easier, since commit 1 is just some polishing. It is commit 2 that adds the e2e example
## Details: Motivation and Example Setup
CoreML partitioner is intended for non-resident models to benefit from CoreML as much as possible. Concretely, if a model cannot be fully converted to CoreML, then CoreML partitioner would help to have the convertible part run on CoreML backend. Unfortunately, for the models we have tried, we cannot find a perfect "torch exportable but CoreML unconvertible" example: they are either fully CoreML resident, or torch unexportable.
As a result, here we artificially create an example by utilizing the `skip_ops_for_coreml_delegation`. It is still a realistic example:
1. In many use cases, we may not want the initial (or final) embedding part to run on CoreML, especially if they are small
2. Our intended graph break emerges from this example
## Appendix: Exported Program
Use `CoreMLPartitioner()` (real user usage)
```
def forward(self, arg152_1: "f32[1, 3, 224, 224]"):
# No stacktrace found for following nodes
lowered_module_0 = self.lowered_module_0
executorch_call_delegate = torch.ops.higher_order.executorch_call_delegate(lowered_module_0, arg152_1); lowered_module_0 = arg152_1 = None
getitem: "f32[1, 1000]" = executorch_call_delegate[0]; executorch_call_delegate = None
return (getitem,)
```
Use `CoreMLPartitioner(skip_ops_for_coreml_delegation=["aten.convolution.default"])` (for demo purpose)
```
def forward(self, arg2_1: "f32[768, 3, 16, 16]", arg3_1: "f32[768]", arg152_1: "f32[1, 3, 224, 224]"):
# File: /Volumes/data/Software/Mine/executorch_121559468/env/lib/python3.10/site-packages/torchvision/models/vision_transformer.py:277 in _process_input, code: x = self.conv_proj(x)
aten_convolution_default: "f32[1, 768, 14, 14]" = executorch_exir_dialects_edge__ops_aten_convolution_default(arg152_1, arg2_1, arg3_1, [16, 16], [0, 0], [1, 1], False, [0, 0], 1); arg152_1 = arg2_1 = arg3_1 = None
# No stacktrace found for following nodes
lowered_module_0 = self.lowered_module_0
executorch_call_delegate = torch.ops.higher_order.executorch_call_delegate(lowered_module_0, aten_convolution_default); lowered_module_0 = aten_convolution_default = None
getitem: "f32[1, 1000]" = executorch_call_delegate[0]; executorch_call_delegate = None
return (getitem,)
```
Pull Request resolved: https://github.com/pytorch/executorch/pull/1981
Reviewed By: cccclai
Differential Revision: D53794713
Pulled By: shoumikhin
fbshipit-source-id: 8d97e7e55fe2f7b787b5ff585e55de5cd391e5c3 Y
yifan_shen3 committed
bdde144cea55ed173b9aa73829d29e15b900fa4a
Parent: 1d07fc4
Committed by Facebook GitHub Bot <facebook-github-bot@users.noreply.github.com>
on 2/15/2024, 7:20:07 AM