Add torchvision_vit model to the examples (#33)

Summary:
Pull Request resolved: #33

Info about the the model: https://pytorch.org/vision/main/models/generated/torchvision.models.vit_b_16.html#torchvision.models.vit_b_16

Differential Revision: D48012005

fbshipit-source-id: 3abf172719af31f97e7f8de4022f97f986e74b5f

Author: guangy10

docs/website/docs/tutorials/00_setting_up_executorch.md

@@ -24,7 +24,7 @@ Setting up PyTorch
 # Install the nightly builds
 # Note: if you are behind a firewall an appropriate proxy server must be setup
 # for all subsequent steps.
-TORCH_VERSION=2.1.0.dev20230809
+TORCH_VERSION=2.1.0.dev20230810
 pip install --force-reinstall --pre torch=="${TORCH_VERSION}" -i https://download.pytorch.org/whl/nightly/cpu
 ```
 

examples/export/test/test_export.py

@@ -11,11 +11,6 @@
 from executorch.examples.export.utils import _CAPTURE_CONFIG, _EDGE_COMPILE_CONFIG
 from executorch.examples.models import MODEL_NAME_TO_MODEL
 
-# pyre-ignore[21]: Could not find module `executorch.extension.pybindings.portable`.
-from executorch.extension.pybindings.portable import (  # @manual
-    _load_for_executorch_from_buffer,
-)
-
 
 class ExportTest(unittest.TestCase):
     def _assert_eager_lowered_same_result(
@@ -28,13 +23,11 @@ def _assert_eager_lowered_same_result(
         )
 
         executorch_model = edge_model.to_executorch()
-        # pyre-ignore
-        pte_model = _load_for_executorch_from_buffer(executorch_model.buffer)
 
         with torch.no_grad():
             eager_output = eager_model(*example_inputs)
         with torch.no_grad():
-            executorch_output = pte_model.forward(example_inputs)
+            executorch_output = executorch_model.exported_program(*example_inputs)
         self.assertTrue(
             torch.allclose(eager_output, executorch_output[0], rtol=1e-5, atol=1e-5)
         )
@@ -50,3 +43,9 @@ def test_mv2_export_to_executorch(self):
         eager_model = eager_model.eval()
 
         self._assert_eager_lowered_same_result(eager_model, example_inputs)
+
+    def test_vit_export_to_executorch(self):
+        eager_model, example_inputs = MODEL_NAME_TO_MODEL["vit"]()
+        eager_model = eager_model.eval()
+
+        self._assert_eager_lowered_same_result(eager_model, example_inputs)

examples/install_requirements.sh

@@ -16,5 +16,5 @@
 # `TORCH_VISION_VERSION` value in this document will be the correct version for the
 # corresponsing version of the repo.
 
-TORCH_VISION_VERSION=0.16.0.dev20230809
+TORCH_VISION_VERSION=0.16.0.dev20230810
 pip install --force-reinstall --pre torchvision=="${TORCH_VISION_VERSION}" -i https://download.pytorch.org/whl/nightly/cpu

examples/models/TARGETS

@@ -10,6 +10,7 @@ python_library(
         "//caffe2:torch",
         "//executorch/examples/models/mobilenet_v2:mv2_export",
         "//executorch/examples/models/mobilenet_v3:mv3_export",
+        "//executorch/examples/models/torchvision_vit:vit_export",
         "//executorch/exir/backend:compile_spec_schema",
     ],
 )

examples/models/models.py

@@ -89,11 +89,18 @@ def gen_mobilenet_v2_model_inputs() -> Tuple[torch.nn.Module, Any]:
     return MV2Model.get_model(), MV2Model.get_example_inputs()
 
 
+def gen_torchvision_vit_model_and_inputs() -> Tuple[torch.nn.Module, Any]:
+    from ..models.torchvision_vit import TorchVisionViTModel
+
+    return TorchVisionViTModel.get_model(), TorchVisionViTModel.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,
+    "vit": gen_torchvision_vit_model_and_inputs,
 }

examples/models/torchvision_vit/TARGETS

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

examples/models/torchvision_vit/__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 .export import TorchVisionViTModel
+
+__all__ = [
+    TorchVisionViTModel,
+]

examples/models/torchvision_vit/export.py

@@ -0,0 +1,30 @@
+# 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.
+
+import logging
+
+import torch
+from torchvision import models
+
+FORMAT = "[%(filename)s:%(lineno)s] %(message)s"
+logging.basicConfig(format=FORMAT)
+
+
+class TorchVisionViTModel:
+    def __init__(self):
+        pass
+
+    @staticmethod
+    def get_model():
+        logging.info("loading torchvision vit_b_16 model")
+        vit_b_16 = models.vit_b_16(weights="IMAGENET1K_V1")
+        logging.info("loaded torchvision vit_b_16 model")
+        return vit_b_16
+
+    @staticmethod
+    def get_example_inputs():
+        input_shape = (1, 3, 224, 224)
+        return (torch.randn(input_shape),)