Export emformer RNNT encode, predict, join (#173)

Summary:
Pull Request resolved: #173

An example to check the export path for Emformer-RNN-T encode, predict, join methods.

Reviewed By: kimishpatel

Differential Revision: D48327041

fbshipit-source-id: 97bb3ee3feadc01ac4f96b5e5702adb867aa2877

Author: kirklandsign

.ci/scripts/gather_test_models.py

@@ -19,7 +19,9 @@
 DEFAULT_RUNNER = "linux.2xlarge"
 RUNNERS = {
     # This one runs OOM on smaller runner, the root cause is unclear (T163016365)
-    "w2l": "linux.12xlarge"
+    "w2l": "linux.12xlarge",
+    # This one causes timeout on smaller runner, the root cause is unclear (T161064121)
+    "emformer_join": "linux.12xlarge",
 }
 
 

examples/export/utils.py

@@ -18,9 +18,9 @@
 
 _CAPTURE_CONFIG = exir.CaptureConfig(enable_aot=True)
 
-# Explicitly force the activation of the IR validator
+# TODO(T163721729): Enable IR check after decomposing div.Tensor_mode
 _EDGE_COMPILE_CONFIG = exir.EdgeCompileConfig(
-    _check_ir_validity=True,
+    _check_ir_validity=False,
 )
 
 

examples/models/TARGETS

@@ -10,6 +10,7 @@ python_library(
         "//caffe2:torch",
         "//executorch/examples/models:model_base",  # @manual
         "//executorch/examples/models/deeplab_v3:dl3_model",  # @manual
+        "//executorch/examples/models/emformer_rnnt:emformer_rnnt_model",  # @manual
         "//executorch/examples/models/inception_v3:ic3_model",  # @manual
         "//executorch/examples/models/inception_v4:ic4_model",  # @manual
         "//executorch/examples/models/mobilebert:mobilebert_model",  # @manual

examples/models/__init__.py

@@ -12,6 +12,9 @@
     "add": ("toy_model", "AddModule"),
     "add_mul": ("toy_model", "AddMulModule"),
     "dl3": ("deeplab_v3", "DeepLabV3ResNet50Model"),
+    "emformer_transcribe": ("emformer_rnnt", "EmformerRnntTranscriberModel"),
+    "emformer_predict": ("emformer_rnnt", "EmformerRnntPredictorModel"),
+    "emformer_join": ("emformer_rnnt", "EmformerRnntJoinerModel"),
     "mobilebert": ("mobilebert", "MobileBertModelExample"),
     "mv2": ("mobilenet_v2", "MV2Model"),
     "mv3": ("mobilenet_v3", "MV3Model"),

examples/models/emformer_rnnt/__init__.py

@@ -0,0 +1,17 @@
+# 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 (
+    EmformerRnntJoinerModel,
+    EmformerRnntPredictorModel,
+    EmformerRnntTranscriberModel,
+)
+
+__all__ = [
+    EmformerRnntTranscriberModel,
+    EmformerRnntPredictorModel,
+    EmformerRnntJoinerModel,
+]

examples/models/emformer_rnnt/model.py

@@ -0,0 +1,132 @@
+# 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 torchaudio
+
+from ..model_base import EagerModelBase
+
+
+FORMAT = "[%(filename)s:%(lineno)s] %(message)s"
+logging.basicConfig(format=FORMAT)
+
+
+__all__ = [
+    "EmformerRnntTranscriberModel",
+    "EmformerRnntPredictorModel",
+    "EmformerRnntJoinerModel",
+]
+
+
+class EmformerRnntTranscriberExample(torch.nn.Module):
+    """
+    This is a wrapper for validating transcriber for the Emformer RNN-T architecture.
+    It does not reflect the actual usage such as beam search, but rather an example for the export workflow.
+    """
+
+    def __init__(self) -> None:
+        super().__init__()
+        bundle = torchaudio.pipelines.EMFORMER_RNNT_BASE_LIBRISPEECH
+        decoder = bundle.get_decoder()
+        m = decoder.model
+        self.rnnt = m
+
+    def forward(self, transcribe_inputs):
+        return self.rnnt.transcribe(*transcribe_inputs)
+
+
+class EmformerRnntTranscriberModel(EagerModelBase):
+    def __init__(self):
+        pass
+
+    def get_eager_model(self) -> torch.nn.Module:
+        logging.info("Loading emformer rnnt transcriber")
+        m = EmformerRnntTranscriberExample()
+        logging.info("Loaded emformer rnnt transcriber")
+        return m
+
+    def get_example_inputs(self):
+        transcribe_inputs = (
+            torch.randn(1, 128, 80),
+            torch.tensor([128]),
+        )
+        return (transcribe_inputs,)
+
+
+class EmformerRnntPredictorExample(torch.nn.Module):
+    """
+    This is a wrapper for validating predictor for the Emformer RNN-T architecture.
+    It does not reflect the actual usage such as beam search, but rather an example for the export workflow.
+    """
+
+    def __init__(self) -> None:
+        super().__init__()
+        bundle = torchaudio.pipelines.EMFORMER_RNNT_BASE_LIBRISPEECH
+        decoder = bundle.get_decoder()
+        m = decoder.model
+        self.rnnt = m
+
+    def forward(self, predict_inputs):
+        return self.rnnt.predict(*predict_inputs)
+
+
+class EmformerRnntPredictorModel(EagerModelBase):
+    def __init__(self):
+        pass
+
+    def get_eager_model(self) -> torch.nn.Module:
+        logging.info("Loading emformer rnnt predictor")
+        m = EmformerRnntPredictorExample()
+        logging.info("Loaded emformer rnnt predictor")
+        return m
+
+    def get_example_inputs(self):
+        predict_inputs = (
+            torch.zeros([1, 128], dtype=int),
+            torch.tensor([128], dtype=int),
+            None,
+        )
+        return (predict_inputs,)
+
+
+class EmformerRnntJoinerExample(torch.nn.Module):
+    """
+    This is a wrapper for validating joiner for the Emformer RNN-T architecture.
+    It does not reflect the actual usage such as beam search, but rather an example for the export workflow.
+    """
+
+    def __init__(self) -> None:
+        super().__init__()
+        bundle = torchaudio.pipelines.EMFORMER_RNNT_BASE_LIBRISPEECH
+        decoder = bundle.get_decoder()
+        m = decoder.model
+        self.rnnt = m
+
+    def forward(self, predict_inputs):
+        return self.rnnt.join(*predict_inputs)
+
+
+class EmformerRnntJoinerModel(EagerModelBase):
+    def __init__(self):
+        pass
+
+    def get_eager_model(self) -> torch.nn.Module:
+        logging.info("Loading emformer rnnt joiner")
+        m = EmformerRnntJoinerExample()
+        logging.info("Loaded emformer rnnt joiner")
+        return m
+
+    def get_example_inputs(self):
+        join_inputs = (
+            torch.rand([1, 128, 1024]),
+            torch.tensor([128]),
+            torch.rand([1, 128, 1024]),
+            torch.tensor([128]),
+        )
+        return (join_inputs,)