Add MicroSpeechLSTM to test models

Author: skywall

backends/nxp/tests/exported_program_vizualize.py

@@ -102,7 +102,7 @@ def name_color(string):  # pseudo-randomization function
             label = ""
             if "val" in node.meta:
                 tensor = node.meta["val"]
-                if isinstance(tensor, tuple):
+                if isinstance(tensor, tuple) or isinstance(tensor, list):
                     tensor = tensor[0]  # Fake tensor
                 label = f"  ({list(tensor.shape)} | {tensor.dtype})"
 

examples/nxp/aot_neutron_compile.py

@@ -23,6 +23,7 @@
 from executorch.examples.models.model_factory import EagerModelFactory
 from executorch.examples.nxp.cifar_net.cifar_net import CifarNet
 from executorch.examples.nxp.cifar_net.cifar_net import test_cifarnet_model
+from executorch.examples.nxp.models.microspeech_lstm.microspeech_lstm import MicroSpeechLSTM
 from executorch.examples.nxp.models.mlperf_tiny import (AnomalyDetection, KeywordSpotting, ImageClassification,
                                                         VisualWakeWords)
 from executorch.exir import ExecutorchBackendConfig
@@ -90,6 +91,7 @@ def get_model_and_inputs_from_name(model_name: str):
     "keyword_spotting": KeywordSpotting,
     "image_classification": ImageClassification,
     "anomaly_detection": AnomalyDetection,
+    "microspeech_lstm": MicroSpeechLSTM,
 }
 
 

examples/nxp/models/microspeech_lstm/__init__.py

@@ -0,0 +1,64 @@
+# Copyright 2025 NXP
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+import itertools
+import os
+from typing import Iterator
+
+import torch
+from torch import nn
+from torch.utils.data import DataLoader
+
+from executorch.examples.models import model_base
+
+
+class MicroSpeechLSTM(model_base.EagerModelBase):
+
+    def __init__(self):
+        self._weights_file = os.path.join(os.path.dirname(__file__), 'microspeech_lstm_best.pth')
+        calibration_dataset_path = os.path.join(os.path.dirname(__file__), 'calibration_data.pt')
+        self._calibration_dataset = torch.load(calibration_dataset_path, map_location=torch.device('cpu'))
+
+    @staticmethod
+    def _collate_fn(data: list[tuple]):
+        data, labels = zip(*data)
+        return torch.stack(list(data)), torch.tensor(list(labels))
+
+    def get_eager_model(self) -> torch.nn.Module:
+        lstm_module = MicroSpeechLSTMModule()
+        lstm_module.load_state_dict(torch.load(self._weights_file, weights_only=True, map_location=torch.device('cpu')))
+
+        return lstm_module
+
+    def get_example_inputs(self) -> tuple[torch.Tensor]:
+        sample = self._calibration_dataset[0]
+
+        return (sample[0].unsqueeze(0),)  # use only sample data, not label
+
+    def get_calibration_inputs(self, batch_size: int = 1) -> Iterator[tuple[torch.Tensor]]:
+        """
+        Get LSTM's calibration input. Batch size is ignored and set to 1 because hidden state
+        has to be initialized to the size of batch.
+
+        :param batch_size: Ignored.
+        :return: Iterator with input calibration dataset samples.
+        """
+        data_loader = DataLoader(self._calibration_dataset, batch_size=1)
+        return itertools.starmap(lambda data, label: (data,), iter(data_loader))
+
+
+class MicroSpeechLSTMModule(nn.Module):
+
+    def __init__(self, input_size=80, output_size=3, features_size=128):
+        super(MicroSpeechLSTMModule, self).__init__()
+        self.lstm = nn.LSTM(input_size, features_size, batch_first=True)
+        self.linear = nn.Linear(features_size, output_size)
+
+    def forward(self, x):
+        _, (x, _) = self.lstm(x)
+        x = x.squeeze(0)
+        x = self.linear(x)
+        x = nn.functional.softmax(x, dim=1)
+        return x

examples/nxp/models/microspeech_lstm/calibration_data.pt

@@ -0,0 +1 @@
+3ffab80556d0e543876479d4de245d76  microspeech_lstm_best.pth