Adding LSTM to example models (#5931)

Summary:
Addition of LSTM model to example models.

Thanks!

Pull Request resolved: #5931

Reviewed By: mergennachin

Differential Revision: D64047539

Pulled By: digantdesai

fbshipit-source-id: 6e0d879027f2b76c482a92b82e05458814efbfeb

Author: SaoirseARM

examples/models/__init__.py

@@ -1,5 +1,6 @@
 # Copyright (c) Meta Platforms, Inc. and affiliates.
 # All rights reserved.
+# Copyright 2024 Arm Limited and/or its affiliates.
 #
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
@@ -16,6 +17,7 @@
     "emformer_predict": ("emformer_rnnt", "EmformerRnntPredictorModel"),
     "emformer_join": ("emformer_rnnt", "EmformerRnntJoinerModel"),
     "llama2": ("llama2", "Llama2Model"),
+    "lstm": ("lstm", "LSTMModel"),
     "mobilebert": ("mobilebert", "MobileBertModelExample"),
     "mv2": ("mobilenet_v2", "MV2Model"),
     "mv2_untrained": ("mobilenet_v2", "MV2UntrainedModel"),

examples/models/lstm/__init__.py

@@ -0,0 +1,12 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+# Copyright 2024 Arm Limited and/or its affiliates.
+#
+# 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 LSTMModel
+
+__all__ = [
+    LSTMModel,
+]

examples/models/lstm/model.py

@@ -0,0 +1,31 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+# Copyright 2024 Arm Limited and/or its affiliates.
+#
+# 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 torch.nn.quantizable.modules import rnn
+
+from ..model_base import EagerModelBase
+
+
+class LSTMModel(EagerModelBase):
+    def __init__(self):
+        pass
+
+    def get_eager_model(self) -> torch.nn.Module:
+        logging.info("Loading LSTM model")
+        lstm = rnn.LSTM(10, 20, 2)
+        logging.info("Loaded LSTM model")
+        return lstm
+
+    def get_example_inputs(self):
+        input_tensor = torch.randn(5, 3, 10)
+        h0 = torch.randn(2, 3, 20)
+        c0 = torch.randn(2, 3, 20)
+        return (input_tensor, (h0, c0))