Add more models as part of GA models

Summary:

I realized that we don't have object detection models in our repo.

This is based on looking at list in https://mlcommons.org/benchmarks/inference-mobile/ and https://ai-benchmark.com/tests.html and coming up with a list that we don't have in our examples.

Author: mergennachin

.github/workflows/trunk.yml

@@ -63,7 +63,7 @@ jobs:
       contents: read
     strategy:
       matrix:
-        model: [linear, add, add_mul, ic3, ic4, mv2, mv3, resnet18, resnet50, vit, w2l, mobilebert, emformer_join, emformer_transcribe]
+        model: [linear, add, add_mul, ic3, ic4, mv2, mv3, resnet18, resnet50, vit, w2l, mobilebert, emformer_join, emformer_transcribe, efficientnet_b4, yolov4_tiny, unet, albert, ssd_mobilenetv2, deeplabv3_mobilenet, bert, bilstm, esrgan, srgan, crnn]
         backend: [portable, xnnpack-quantization-delegation]
         runner: [linux.arm64.2xlarge]
         include:

examples/models/__init__.py

@@ -37,6 +37,18 @@ class Model(str, Enum):
     EfficientSam = "efficient_sam"
     Qwen25 = "qwen2_5"
     Phi4Mini = "phi_4_mini"
+    # Newly added models
+    EfficientNetB4 = "efficientnet_b4"
+    YOLOv4Tiny = "yolov4_tiny"
+    UNet = "unet"
+    Albert = "albert"
+    SSDMobileNetV2 = "ssd_mobilenetv2"
+    DeepLabV3MobileNet = "deeplabv3_mobilenet"
+    Bert = "bert"
+    BiLSTM = "bilstm"
+    ESRGAN = "esrgan"
+    SRGAN = "srgan"
+    CRNN = "crnn"
 
     def __str__(self) -> str:
         return self.value
@@ -82,6 +94,18 @@ def __str__(self) -> str:
     str(Model.EfficientSam): ("efficient_sam", "EfficientSAM"),
     str(Model.Qwen25): ("qwen2_5", "Qwen2_5Model"),
     str(Model.Phi4Mini): ("phi_4_mini", "Phi4MiniModel"),
+    # Newly added models
+    str(Model.EfficientNetB4): ("efficientnet_b4", "EfficientNetB4Model"),
+    str(Model.YOLOv4Tiny): ("yolov4_tiny", "YOLOv4TinyModel"),
+    str(Model.UNet): ("unet", "UNetBinaryModel"),
+    str(Model.Albert): ("albert", "AlbertModelExample"),
+    str(Model.SSDMobileNetV2): ("ssd_mobilenetv2", "SSDMobileNetV2Model"),
+    str(Model.DeepLabV3MobileNet): ("deeplabv3_mobilenet", "DeepLabV3MobileNetModel"),
+    str(Model.Bert): ("bert", "BertModelExample"),
+    str(Model.BiLSTM): ("bilstm", "BiLSTMModel"),
+    str(Model.ESRGAN): ("esrgan", "ESRGANModel"),
+    str(Model.SRGAN): ("srgan", "SRGANModel"),
+    str(Model.CRNN): ("crnn", "CRNNModel"),
 }
 
 __all__ = [

examples/models/albert/__init__.py

@@ -0,0 +1,12 @@
+# 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 .model import AlbertLargeModelExample, AlbertModelExample
+
+__all__ = [
+    "AlbertLargeModelExample",
+    "AlbertModelExample",
+]

examples/models/albert/model.py

@@ -0,0 +1,47 @@
+# 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 transformers import AlbertModel, AutoTokenizer  # @manual
+
+from ..model_base import EagerModelBase
+
+
+class AlbertModelExample(EagerModelBase):
+    def __init__(self):
+        pass
+
+    def get_eager_model(self) -> torch.nn.Module:
+        logging.info("Loading ALBERT model")
+        # pyre-ignore
+        model = AlbertModel.from_pretrained("albert-base-v2", return_dict=False)
+        model.eval()
+        logging.info("Loaded ALBERT model")
+        return model
+
+    def get_example_inputs(self):
+        tokenizer = AutoTokenizer.from_pretrained("albert-base-v2")
+        return (tokenizer("Hello, my dog is cute", return_tensors="pt")["input_ids"],)
+
+
+class AlbertLargeModelExample(EagerModelBase):
+    def __init__(self):
+        pass
+
+    def get_eager_model(self) -> torch.nn.Module:
+        logging.info("Loading ALBERT-Large model")
+        # pyre-ignore
+        model = AlbertModel.from_pretrained("albert-large-v2", return_dict=False)
+        model.eval()
+        logging.info("Loaded ALBERT-Large model")
+        return model
+
+    def get_example_inputs(self):
+        tokenizer = AutoTokenizer.from_pretrained("albert-large-v2")
+        return (tokenizer("Hello, my dog is cute", return_tensors="pt")["input_ids"],)

examples/models/bert/__init__.py

@@ -0,0 +1,12 @@
+# 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 .model import BertLargeModelExample, BertModelExample
+
+__all__ = [
+    "BertLargeModelExample",
+    "BertModelExample",
+]

examples/models/bert/model.py

@@ -0,0 +1,51 @@
+# 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 transformers import AutoTokenizer, BertModel  # @manual
+
+from ..model_base import EagerModelBase
+
+
+class BertModelExample(EagerModelBase):
+    def __init__(self):
+        pass
+
+    def get_eager_model(self) -> torch.nn.Module:
+        logging.info("Loading BERT model")
+        # pyre-ignore
+        model = BertModel.from_pretrained(
+            "google-bert/bert-base-uncased", return_dict=False
+        )
+        model.eval()
+        logging.info("Loaded BERT model")
+        return model
+
+    def get_example_inputs(self):
+        tokenizer = AutoTokenizer.from_pretrained("google-bert/bert-base-uncased")
+        return (tokenizer("Hello, my dog is cute", return_tensors="pt")["input_ids"],)
+
+
+class BertLargeModelExample(EagerModelBase):
+    def __init__(self):
+        pass
+
+    def get_eager_model(self) -> torch.nn.Module:
+        logging.info("Loading BERT-Large model")
+        # pyre-ignore
+        model = BertModel.from_pretrained(
+            "google-bert/bert-large-uncased", return_dict=False
+        )
+        model.eval()
+        logging.info("Loaded BERT-Large model")
+        return model
+
+    def get_example_inputs(self):
+        tokenizer = AutoTokenizer.from_pretrained("google-bert/bert-large-uncased")
+        return (tokenizer("Hello, my dog is cute", return_tensors="pt")["input_ids"],)

examples/models/bilstm/__init__.py

@@ -0,0 +1,12 @@
+# 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 .model import BidirectionalLSTMModel, BidirectionalLSTMTextModel
+
+__all__ = [
+    "BidirectionalLSTMModel",
+    "BidirectionalLSTMTextModel",
+]

examples/models/bilstm/model.py

@@ -0,0 +1,115 @@
+# 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
+import torch.nn as nn
+
+from ..model_base import EagerModelBase
+
+
+class BidirectionalLSTM(nn.Module):
+    """Bidirectional LSTM for sequence modeling"""
+
+    def __init__(self, input_size=100, hidden_size=128, num_layers=2, num_classes=10):
+        super(BidirectionalLSTM, self).__init__()
+        self.hidden_size = hidden_size
+        self.num_layers = num_layers
+
+        # Bidirectional LSTM
+        self.lstm = nn.LSTM(
+            input_size, hidden_size, num_layers, batch_first=True, bidirectional=True
+        )
+
+        # Output layer (hidden_size * 2 because of bidirectional)
+        self.fc = nn.Linear(hidden_size * 2, num_classes)
+
+    def forward(self, x):
+        # Initialize hidden states
+        # For bidirectional: hidden states shape is (num_layers * 2, batch, hidden_size)
+        h0 = torch.zeros(self.num_layers * 2, x.size(0), self.hidden_size).to(x.device)
+        c0 = torch.zeros(self.num_layers * 2, x.size(0), self.hidden_size).to(x.device)
+
+        # LSTM forward pass
+        out, _ = self.lstm(x, (h0, c0))
+
+        # Take the last time step output
+        out = self.fc(out[:, -1, :])
+        return out
+
+
+class BidirectionalLSTMTextClassifier(nn.Module):
+    """Bidirectional LSTM for text classification with embedding layer"""
+
+    def __init__(
+        self, vocab_size=10000, embedding_dim=128, hidden_size=256, num_classes=2
+    ):
+        super(BidirectionalLSTMTextClassifier, self).__init__()
+        self.hidden_size = hidden_size
+
+        # Embedding layer
+        self.embedding = nn.Embedding(vocab_size, embedding_dim)
+
+        # Bidirectional LSTM
+        self.lstm = nn.LSTM(
+            embedding_dim, hidden_size, bidirectional=True, batch_first=True
+        )
+
+        # Output layer
+        self.fc = nn.Linear(hidden_size * 2, num_classes)
+
+    def forward(self, x):
+        # Embedding
+        embedded = self.embedding(x)
+
+        # LSTM
+        lstm_out, _ = self.lstm(embedded)
+
+        # Global max pooling over sequence dimension
+        pooled = torch.max(lstm_out, dim=1)[0]
+
+        # Classification
+        output = self.fc(pooled)
+        return output
+
+
+class BidirectionalLSTMModel(EagerModelBase):
+    def __init__(self):
+        pass
+
+    def get_eager_model(self) -> torch.nn.Module:
+        logging.info("Loading Bidirectional LSTM model")
+        model = BidirectionalLSTM(
+            input_size=100, hidden_size=128, num_layers=2, num_classes=10
+        )
+        model.eval()
+        logging.info("Loaded Bidirectional LSTM model")
+        return model
+
+    def get_example_inputs(self):
+        # Example: (batch_size=1, seq_len=50, input_size=100)
+        tensor_size = (1, 50, 100)
+        return (torch.randn(tensor_size),)
+
+
+class BidirectionalLSTMTextModel(EagerModelBase):
+    def __init__(self):
+        pass
+
+    def get_eager_model(self) -> torch.nn.Module:
+        logging.info("Loading Bidirectional LSTM text classifier")
+        model = BidirectionalLSTMTextClassifier(
+            vocab_size=10000, embedding_dim=128, hidden_size=256, num_classes=2
+        )
+        model.eval()
+        logging.info("Loaded Bidirectional LSTM text classifier")
+        return model
+
+    def get_example_inputs(self):
+        # Example: (batch_size=1, seq_len=100) - token indices
+        tensor_size = (1, 100)
+        return (torch.randint(0, 10000, tensor_size),)

examples/models/crnn/__init__.py

@@ -0,0 +1,13 @@
+# 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 .model import CRNNMobileModel, CRNNModel, CRNNRGBModel
+
+__all__ = [
+    "CRNNMobileModel",
+    "CRNNModel",
+    "CRNNRGBModel",
+]