Add model examples and export scripts

Summary:
This diff adds similar baseclass as model inventory;s eagermodelbase.

MV3 example inherits from it and extract the model from torchvision.

Class supports extracting to_edge variant of the model which can be used for
delegation.

Reviewed By: mergennachin

Differential Revision: D47635308

fbshipit-source-id: 43a6a341e19f474e83b22f033bd06b21a003d6ab

Author: kimishpatel

examples/README.md

@@ -0,0 +1,9 @@
+## Dependencies
+
+Various models listed in this directory have dependencies on some other packages, e.g. torchvision, torchaudio.
+In order to make sure model's listed in examples are importable, e.g. via
+```
+from executorch.examples.models.mobilenet_v3d import MV3Model
+m = MV3Model.get_model()
+```
+we need to have appropriate packages installed. You should install these deps via install_requirements.sh

examples/TARGETS

@@ -0,0 +1,11 @@
+load("@fbcode_macros//build_defs:python_library.bzl", "python_library")
+
+python_library(
+    name = "utils",
+    srcs = [
+        "utils.py",
+    ],
+    deps = [
+        "//executorch/exir:lib",
+    ],
+)

examples/export/export_example.py

@@ -5,6 +5,8 @@
 import executorch.exir as exir
 
 import torch
+from executorch.examples.models.mobilenet_v3 import MV3Model
+from executorch.examples.utils import _CAPTURE_CONFIG, _EDGE_COMPILE_CONFIG
 
 
 class MulModule(torch.nn.Module):
@@ -14,7 +16,7 @@ def __init__(self) -> None:
     def forward(self, input, other):
         return input * other
 
-    def get_random_inputs(self):
+    def get_example_inputs(self):
         return (torch.randn(3, 2), torch.randn(3, 2))
 
 
@@ -26,7 +28,7 @@ def __init__(self):
     def forward(self, arg):
         return self.linear(arg)
 
-    def get_random_inputs(self):
+    def get_example_inputs(self):
         return (torch.randn(3, 3),)
 
 
@@ -41,22 +43,23 @@ def forward(self, x, y):
         z = z + z
         return z
 
-    def get_random_inputs(self):
+    def get_example_inputs(self):
         return (torch.ones(1), torch.ones(1))
 
 
 MODEL_NAME_TO_MODEL = {
     "mul": MulModule,
     "linear": LinearModule,
     "add": AddModule,
+    "mv3": None,
 }
 
 
-def export_to_ff(model_name, model):
-    m = model()
-    edge = exir.capture(
-        m, m.get_random_inputs(), exir.CaptureConfig(enable_dynamic_shape=True)
-    ).to_edge(exir.EdgeCompileConfig(_check_ir_validity=False, _use_edge_ops=True))
+def export_to_ff(model_name, model, example_inputs):
+    m = model
+    edge = exir.capture(m, example_inputs, _CAPTURE_CONFIG).to_edge(
+        _EDGE_COMPILE_CONFIG
+    )
     print("Exported graph:\n", edge.graph)
 
     exec_prog = edge.to_executorch()
@@ -85,6 +88,17 @@ def export_to_ff(model_name, model):
             f"Model {args.model_name} is not a valid name. "
             f"Available models are {list(MODEL_NAME_TO_MODEL.keys())}."
         )
-    model = MODEL_NAME_TO_MODEL[args.model_name]
 
-    export_to_ff(args.model_name, model)
+    if args.model_name == "mv3":
+        # Unfortunately lack of consistent interface on example models in this file
+        # and how we obtain oss models result in changes like this.
+        # we should probably fix this if all the MVP model's export example
+        # wiil be added here.
+        # For now, to unblock, not planning to land those changes in the current diff
+        model = MV3Model.get_model().eval()
+        example_inputs = MV3Model.get_example_inputs()
+    else:
+        model = MODEL_NAME_TO_MODEL[args.model_name]()
+        example_inputs = model.get_example_inputs()
+
+    export_to_ff(args.model_name, model, example_inputs)

examples/install_requirements.sh

@@ -0,0 +1,7 @@
+#!/bin/bash
+
+# install pre-requisite
+# here it is used to install torchvision's nighlty package because the latest
+# variant install an older version of pytorch, 1.8,
+# tested only on linux
+pip install --pre torchvision -i https://download.pytorch.org/whl/nightly/cpu

examples/models/mobilenet_v3/TARGETS

@@ -0,0 +1,14 @@
+load("@fbcode_macros//build_defs:python_library.bzl", "python_library")
+
+python_library(
+    name = "mv3_export",
+    srcs = [
+        "__init__.py",
+        "export.py",
+    ],
+    base_module = "executorch.examples.models.mobilenet_v3",
+    deps = [
+        "//caffe2:torch",
+        "//pytorch/vision:torchvision",
+    ],
+)

examples/models/mobilenet_v3/__init__.py

@@ -0,0 +1,5 @@
+from executorch.examples.models.mobilenet_v3.export import MV3Model
+
+__all__ = [
+    MV3Model,
+]

examples/models/mobilenet_v3/export.py

@@ -0,0 +1,24 @@
+import logging
+
+import torch
+from torchvision import models
+
+FORMAT = "[%(filename)s:%(lineno)s] %(message)s"
+logging.basicConfig(format=FORMAT)
+
+# will refactor this in a separate file.
+class MV3Model:
+    def __init__(self):
+        pass
+
+    @staticmethod
+    def get_model():
+        logging.info("loading mobilenet_v3 model")
+        mv3_small = models.mobilenet_v3_small(pretrained=True)
+        logging.info("loaded mobilenet_v3 model")
+        return mv3_small
+
+    @staticmethod
+    def get_example_inputs():
+        tensor_size = (1, 3, 224, 224)
+        return (torch.randn(tensor_size),)

examples/models/mobilenet_v3/test/TARGETS

@@ -0,0 +1,15 @@
+load("@fbcode_macros//build_defs:python_unittest.bzl", "python_unittest")
+
+python_unittest(
+    name = "test_export",
+    srcs = [
+        "test_export.py",
+    ],
+    supports_static_listing = True,
+    deps = [
+        "//caffe2:torch",
+        "//executorch/examples:utils",
+        "//executorch/examples/models/mobilenet_v3:mv3_export",
+        "//executorch/exir:lib",
+    ],
+)

examples/models/mobilenet_v3/test/test_export.py

@@ -0,0 +1,25 @@
+import unittest
+
+import torch
+from executorch.examples.models.mobilenet_v3 import MV3Model
+from executorch.examples.utils import _EDGE_COMPILE_CONFIG
+
+
+class ExportTest(unittest.TestCase):
+    def test_export_to_executorch(self):
+        eager_model = MV3Model.get_model().eval()
+        import executorch.exir as exir
+
+        capture_config = exir.CaptureConfig(enable_dynamic_shape=False)
+        edge_model = exir.capture(
+            eager_model, MV3Model.get_example_inputs(), capture_config
+        ).to_edge(_EDGE_COMPILE_CONFIG)
+        example_inputs = MV3Model.get_example_inputs()
+        with torch.no_grad():
+            eager_output = eager_model(*example_inputs)
+        executorch_model = edge_model.to_executorch()
+        with torch.no_grad():
+            executorch_output = executorch_model.graph_module(*example_inputs)
+        self.assertTrue(
+            torch.allclose(eager_output, executorch_output[0], rtol=1e-5, atol=1e-5)
+        )

examples/utils.py

@@ -0,0 +1,11 @@
+import executorch.exir as exir
+
+# Using dynamic shape does not allow us to run graph_module returned by
+# to_executorch for mobilenet_v3.
+# Reason is that there memory allocation ops with symbolic shape nodes.
+# and when evaulating shape, it doesnt seem that we presenting them with shape env
+# that contain those variables.
+_CAPTURE_CONFIG = exir.CaptureConfig(enable_dynamic_shape=True)
+_EDGE_COMPILE_CONFIG = exir.EdgeCompileConfig(
+    _check_ir_validity=False, _use_edge_ops=True
+)