Put models retrival API under examples/models

Summary: Extract out `MODEL_NAME_TO_MODEL` dict into a separate file and put it under `examples/models`. This can then be reused.

Reviewed By: jerryzh168

Differential Revision: D47885828

fbshipit-source-id: 40ca0508293f97c5c66a25a85e125a4a6c322d61

Author: larryliu0820

examples/export/export_and_delegate.py

@@ -11,10 +11,11 @@
 import executorch.exir as exir
 import torch
 from executorch.backends.backend_api import to_backend
-from executorch.backends.compile_spec_schema import CompileSpec
 from executorch.backends.test.backend_with_compiler_demo import BackendWithCompilerDemo
 from executorch.backends.test.op_partitioner_demo import AddMulPartitionerDemo
 
+from ..models import MODEL_NAME_TO_MODEL
+
 from .utils import _CAPTURE_CONFIG, _EDGE_COMPILE_CONFIG
 
 """
@@ -28,23 +29,6 @@
 """
 
 
-class AddMulModule(torch.nn.Module):
-    def __init__(self):
-        super().__init__()
-
-    def forward(self, a, x, b):
-        y = torch.mm(a, x)
-        z = torch.add(y, b)
-        return z
-
-    def get_random_inputs(self):
-        return (torch.ones(2, 2), 2 * torch.ones(2, 2), 3 * torch.ones(2, 2))
-
-    def get_compile_spec(self):
-        max_value = self.get_random_inputs()[0].shape[0]
-        return [CompileSpec("max_value", bytes([max_value]))]
-
-
 def export_compsite_module_with_lower_graph():
     """
 
@@ -63,9 +47,9 @@ def export_compsite_module_with_lower_graph():
 
     """
     print("Running the example to export a composite module with lowered graph...")
-
-    m = AddMulModule().eval()
-    m_inputs = m.get_random_inputs()
+    m, m_inputs = MODEL_NAME_TO_MODEL.get("add_mul")()
+    m = m.eval()
+    m_inputs = m.get_example_inputs()
     edge = exir.capture(m, m_inputs, _CAPTURE_CONFIG).to_edge(_EDGE_COMPILE_CONFIG)
     print("Exported graph:\n", edge.exported_program.graph)
 
@@ -169,10 +153,10 @@ def export_and_lower_the_whole_graph():
     """
     print("Running the example to export and lower the whole graph...")
 
-    m = AddMulModule()
-    edge = exir.capture(m, m.get_random_inputs(), _CAPTURE_CONFIG).to_edge(
-        _EDGE_COMPILE_CONFIG
-    )
+    m, m_inputs = MODEL_NAME_TO_MODEL.get("add_mul")()
+    m = m.eval()
+    m_inputs = m.get_example_inputs()
+    edge = exir.capture(m, m_inputs, _CAPTURE_CONFIG).to_edge(_EDGE_COMPILE_CONFIG)
     print("Exported graph:\n", edge.exported_program.graph)
 
     # Lower AddMulModule to the demo backend

examples/export/export_example.py

@@ -7,82 +7,14 @@
 # Example script for exporting simple models to flatbuffer
 
 import argparse
-from typing import Any, Tuple
 
 import executorch.exir as exir
 
-import torch
+from ..models import MODEL_NAME_TO_MODEL
 
 from .utils import _CAPTURE_CONFIG, _EDGE_COMPILE_CONFIG
 
 
-class MulModule(torch.nn.Module):
-    def __init__(self) -> None:
-        super().__init__()
-
-    def forward(self, input, other):
-        return input * other
-
-    @staticmethod
-    def get_example_inputs():
-        return (torch.randn(3, 2), torch.randn(3, 2))
-
-
-class LinearModule(torch.nn.Module):
-    def __init__(self):
-        super().__init__()
-        self.linear = torch.nn.Linear(3, 3)
-
-    def forward(self, arg):
-        return self.linear(arg)
-
-    @staticmethod
-    def get_example_inputs():
-        return (torch.randn(3, 3),)
-
-
-class AddModule(torch.nn.Module):
-    def __init__(self):
-        super().__init__()
-
-    def forward(self, x, y):
-        z = x + y
-        z = z + x
-        z = z + x
-        z = z + z
-        return z
-
-    @staticmethod
-    def get_example_inputs():
-        return (torch.ones(1), torch.ones(1))
-
-
-def gen_mobilenet_v3_model_inputs() -> Tuple[torch.nn.Module, Any]:
-    # 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
-    from ..models.mobilenet_v3 import MV3Model
-
-    return MV3Model.get_model(), MV3Model.get_example_inputs()
-
-
-def gen_mobilenet_v2_model_inputs() -> Tuple[torch.nn.Module, Any]:
-    from ..models.mobilenet_v2 import MV2Model
-
-    return MV2Model.get_model(), MV2Model.get_example_inputs()
-
-
-MODEL_NAME_TO_MODEL = {
-    "mul": lambda: (MulModule(), MulModule.get_example_inputs()),
-    "linear": lambda: (LinearModule(), LinearModule.get_example_inputs()),
-    "add": lambda: (AddModule(), AddModule.get_example_inputs()),
-    "mv2": gen_mobilenet_v2_model_inputs,
-    "mv3": gen_mobilenet_v3_model_inputs,
-}
-
-
 def export_to_ff(model_name, model, example_inputs):
     m = model.eval()
     edge = exir.capture(m, example_inputs, _CAPTURE_CONFIG).to_edge(

examples/export/test/TARGETS

@@ -9,8 +9,7 @@ python_unittest(
     deps = [
         "//caffe2:torch",
         "//executorch/examples/export:utils",
-        "//executorch/examples/models/mobilenet_v2:mv2_export",
-        "//executorch/examples/models/mobilenet_v3:mv3_export",
+        "//executorch/examples/models:models",
         "//executorch/exir:lib",
     ],
 )

examples/export/test/test_export.py

@@ -7,10 +7,9 @@
 import unittest
 
 import torch
-from executorch.examples.models.mobilenet_v2 import MV2Model
-from executorch.examples.models.mobilenet_v3 import MV3Model
 
-from ..utils import _EDGE_COMPILE_CONFIG
+from executorch.examples.export.utils import _EDGE_COMPILE_CONFIG
+from executorch.examples.models import MODEL_NAME_TO_MODEL
 
 
 class ExportTest(unittest.TestCase):
@@ -34,13 +33,13 @@ def _assert_eager_lowered_same_result(
         )
 
     def test_mv3_export_to_executorch(self):
-        eager_model = MV3Model.get_model().eval()
-        example_inputs = MV3Model.get_example_inputs()
+        eager_model, example_inputs = MODEL_NAME_TO_MODEL["mv3"]()
+        eager_model = eager_model.eval()
 
         self._assert_eager_lowered_same_result(eager_model, example_inputs)
 
     def test_mv2_export_to_executorch(self):
-        eager_model = MV2Model.get_model().eval()
-        example_inputs = MV2Model.get_example_inputs()
+        eager_model, example_inputs = MODEL_NAME_TO_MODEL["mv2"]()
+        eager_model = eager_model.eval()
 
         self._assert_eager_lowered_same_result(eager_model, example_inputs)

examples/models/TARGETS

@@ -0,0 +1,15 @@
+load("@fbcode_macros//build_defs:python_library.bzl", "python_library")
+
+python_library(
+    name = "models",
+    srcs = [
+        "__init__.py",
+        "models.py",
+    ],
+    deps = [
+        "//caffe2:torch",
+        "//executorch/backends:compile_spec_schema",
+        "//executorch/examples/models/mobilenet_v2:mv2_export",
+        "//executorch/examples/models/mobilenet_v3:mv3_export",
+    ],
+)

examples/models/__init__.py

@@ -0,0 +1,11 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+from .models import MODEL_NAME_TO_MODEL
+
+__all__ = [
+    MODEL_NAME_TO_MODEL,
+]

examples/models/models.py

@@ -0,0 +1,99 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+# @file models.py
+# Simple models for demonstration purposes.
+
+from typing import Any, Tuple
+
+import torch
+from executorch.backends.compile_spec_schema import CompileSpec
+
+
+class MulModule(torch.nn.Module):
+    def __init__(self) -> None:
+        super().__init__()
+
+    def forward(self, input, other):
+        return input * other
+
+    @staticmethod
+    def get_example_inputs():
+        return (torch.randn(3, 2), torch.randn(3, 2))
+
+
+class LinearModule(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.linear = torch.nn.Linear(3, 3)
+
+    def forward(self, arg):
+        return self.linear(arg)
+
+    @staticmethod
+    def get_example_inputs():
+        return (torch.randn(3, 3),)
+
+
+class AddModule(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+
+    def forward(self, x, y):
+        z = x + y
+        z = z + x
+        z = z + x
+        z = z + z
+        return z
+
+    @staticmethod
+    def get_example_inputs():
+        return (torch.ones(1), torch.ones(1))
+
+
+class AddMulModule(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+
+    def forward(self, a, x, b):
+        y = torch.mm(a, x)
+        z = torch.add(y, b)
+        return z
+
+    @staticmethod
+    def get_example_inputs():
+        return (torch.ones(2, 2), 2 * torch.ones(2, 2), 3 * torch.ones(2, 2))
+
+    def get_compile_spec(self):
+        max_value = self.get_random_inputs()[0].shape[0]
+        return [CompileSpec("max_value", bytes([max_value]))]
+
+
+def gen_mobilenet_v3_model_inputs() -> Tuple[torch.nn.Module, Any]:
+    # 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
+    from ..models.mobilenet_v3 import MV3Model
+
+    return MV3Model.get_model(), MV3Model.get_example_inputs()
+
+
+def gen_mobilenet_v2_model_inputs() -> Tuple[torch.nn.Module, Any]:
+    from ..models.mobilenet_v2 import MV2Model
+
+    return MV2Model.get_model(), MV2Model.get_example_inputs()
+
+
+MODEL_NAME_TO_MODEL = {
+    "mul": lambda: (MulModule(), MulModule.get_example_inputs()),
+    "linear": lambda: (LinearModule(), LinearModule.get_example_inputs()),
+    "add": lambda: (AddModule(), AddModule.get_example_inputs()),
+    "add_mul": lambda: (AddMulModule(), AddMulModule.get_example_inputs()),
+    "mv2": gen_mobilenet_v2_model_inputs,
+    "mv3": gen_mobilenet_v3_model_inputs,
+}