Pull request pytorch#71: Modify MLPerf Tiny models to accept datasets with labels

Merge in AITEC/executorch from feature/nxg11066/EIEX-184-Add-labels-to-all-MLPerf-Tiny-calibration-data to main-nxp

* commit '1ffc1f7c47740633cdf7c7a6aff65b5cc46d7d79':
  Modify MLPerf Tiny models to accept datasets with labels

Author: roman-janik-nxp

examples/nxp/README.md

@@ -48,7 +48,7 @@ We support several example models. To run these examples, we will use `aot_neutr
 
     Run the script with following arguments: `--model_name` - specify model from a list of available models, `--so_library` 
     - specify path to `libquantized_ops_aot_lib.so` library from step 1. Quantization and delegation are controlled by
-    `--quantize` and `--quantize` flags, turned off by default.
+    `--quantize` and `--delegate` flags, turned off by default.
 
     Supported models:
     - cifar10

examples/nxp/models/mlperf_tiny/anomaly_detection/anomaly_detection.py

@@ -4,7 +4,6 @@
 # LICENSE file in the root directory of this source tree.
 
 import os
-from typing import Iterator
 
 import torch
 
@@ -27,14 +26,5 @@ def _dts_path(self):
     def _input_shape(self):
         return self.__input_shape
 
-    @staticmethod
-    def _collate_fn(data: torch.Tensor, **kwargs):
-        return (torch.stack(list(data)),)
-
-    def get_calibration_inputs(self, batch_size: int = 1) -> Iterator[tuple[torch.Tensor]]:
-        self._batch_size = batch_size
-        data_loader = self._get_data_loader()
-        return iter(data_loader)
-
     def get_eager_model(self) -> torch.nn.Module:
         return self._model_manager.get_model("anomaly_detection", input_dimension=self._input_dim)

examples/nxp/models/mlperf_tiny/image_classification/image_classification.py

@@ -4,7 +4,6 @@
 # LICENSE file in the root directory of this source tree.
 
 import os
-from typing import Iterator
 
 import torch
 
@@ -26,14 +25,5 @@ def _dts_path(self):
     def _input_shape(self):
         return self.__input_shape
 
-    @staticmethod
-    def _collate_fn(data: torch.Tensor, **kwargs):
-        return (torch.stack(list(data)),)
-
-    def get_calibration_inputs(self, batch_size: int = 1) -> Iterator[tuple[torch.Tensor]]:
-            self._batch_size = batch_size
-            data_loader = self._get_data_loader()
-            return iter(data_loader)
-
     def get_eager_model(self) -> torch.nn.Module:
         return self._model_manager.get_model("image_classification")

examples/nxp/models/mlperf_tiny/keyword_spotting/keyword_spotting.py

@@ -3,9 +3,7 @@
 # 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
 
@@ -27,17 +25,6 @@ def _dts_path(self):
     def _input_shape(self):
         return self.__input_shape
 
-    @staticmethod
-    def _collate_fn(data: torch.Tensor, **kwargs):
-        data, labels = zip(*data)
-        return torch.stack(list(data)).to(memory_format=torch.channels_last), torch.tensor(list(labels))
-
-    def get_calibration_inputs(self, batch_size: int = 1) -> Iterator[tuple[torch.Tensor]]:
-        self._batch_size = batch_size
-        data_loader = self._get_data_loader()
-        get_first = lambda a, b: (a,)
-        return itertools.starmap(get_first, iter(data_loader))
-
     def get_eager_model(self) -> torch.nn.Module:
         example_input_shape = self._input_shape[1:]
         return self._model_manager.get_model("keyword_spotting", input_shape=example_input_shape, num_classes=12)

examples/nxp/models/mlperf_tiny/mlperf_tiny_model.py

@@ -3,8 +3,10 @@
 # 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 abc import abstractmethod
+from typing import Iterator
 
 import torch
 from torch.utils.data import DataLoader
@@ -33,9 +35,16 @@ def _dts_path(self):
     def _input_shape(self):
         pass
 
-    @abstractmethod
-    def _collate_fn(self, data: torch.Tensor):
-        pass
+    @staticmethod
+    def _collate_fn(data: list[tuple]):
+        data, labels = zip(*data)
+        return torch.stack(list(data)), torch.tensor(list(labels))
+
+    def get_calibration_inputs(self, batch_size: int = 1) -> Iterator[tuple[torch.Tensor]]:
+        self._batch_size = batch_size
+        data_loader = self._get_data_loader()
+        get_first = lambda a, b: (a,)
+        return itertools.starmap(get_first, iter(data_loader))
 
     @abstractmethod
     def get_eager_model(self) -> torch.nn.Module:
@@ -53,4 +62,7 @@ def _get_data_loader(self):
 
     def _init_dataset(self):
         if self._dataset is None:
+            os.makedirs(os.path.dirname(self._dts_path), exist_ok=True)
+            if not os.path.exists(self._dts_path):
+                raise FileNotFoundError("Calibration data file not found! For more info, follow README.md.")
             self._dataset = CalibrationDataset(self._dts_path)

examples/nxp/models/mlperf_tiny/visual_wake_words/visual_wake_words.py

@@ -3,9 +3,7 @@
 # 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
 
@@ -27,16 +25,5 @@ def _dts_path(self):
     def _input_shape(self):
         return self.__input_shape
 
-    @staticmethod
-    def _collate_fn(data: torch.Tensor, **kwargs):
-        images, labels = zip(*data)
-        return torch.stack(list(images)).to(memory_format=torch.channels_last), torch.tensor(list(labels))
-
-    def get_calibration_inputs(self, batch_size: int = 1) -> Iterator[tuple[torch.Tensor]]:
-        self._batch_size = batch_size
-        data_loader = self._get_data_loader()
-        get_first = lambda a, b: (a,)
-        return itertools.starmap(get_first, iter(data_loader))
-
     def get_eager_model(self) -> torch.nn.Module:
         return self._model_manager.get_model("visual_wake_words")