fixing ONNX and PyTorch model/data tests (#55)

* ONNX model and data test fixes
* pytorch model test fixes

Author: bfineran

src/sparseml/onnx/optim/sensitivity_pruning.py

@@ -389,7 +389,7 @@ def pruning_perf_sens_one_shot_iter(
     model: Union[str, ModelProto],
     data: DataLoader,
     batch_size: int,
-    num_cores: int = -1,
+    num_cores: int = None,
     iterations_per_check: int = 10,
     warmup_iterations_per_check: int = 5,
     sparsity_levels: List[float] = default_pruning_sparsities_perf(),
@@ -408,7 +408,8 @@ def pruning_perf_sens_one_shot_iter(
         to calculate the sparse sensitivity analysis for
     :param data: the data to run through the model
     :param batch_size: the size of the batch to create the model in neural magic for
-    :param num_cores: number of physical cores to run on
+    :param num_cores: number of physical cores to run on. Default is the maximum number
+        of cores available
     :param iterations_per_check: number of iterations to run for perf details
     :param warmup_iterations_per_check: number of iterations to run before perf details
     :param sparsity_levels: the sparsity levels to calculate the loss for for each param
@@ -484,7 +485,7 @@ def pruning_perf_sens_one_shot(
     model: Union[str, ModelProto],
     data: DataLoader,
     batch_size: int,
-    num_cores: int = -1,
+    num_cores: int = None,
     iterations_per_check: int = 10,
     warmup_iterations_per_check: int = 5,
     sparsity_levels: List[float] = default_pruning_sparsities_perf(),
@@ -500,7 +501,8 @@ def pruning_perf_sens_one_shot(
         to calculate the sparse sensitivity analysis for
     :param data: the data to run through the model
     :param batch_size: the size of the batch to create the model in neural magic for
-    :param num_cores: number of physical cores to run on
+    :param num_cores: number of physical cores to run on. Default is the maximum
+        available
     :param iterations_per_check: number of iterations to run for perf details
     :param warmup_iterations_per_check: number of iterations to run before perf details
     :param sparsity_levels: the sparsity levels to calculate the loss for for each param

src/sparseml/onnx/utils/data.py

@@ -48,8 +48,10 @@ class DataLoader(object):
     Iterator returns a tuple containing (data, label).
     label is only returned if label data was passed in.
 
-    :param data: a file glob pointing to numpy files or loaded numpy data
-    :param labels: a file glob pointing to numpy files or loaded numpy data
+    :param data: a file glob pointing to numpy files, path to a tar ball of numpy
+        files, or loaded numpy data
+    :param labels: a file glob pointing to numpy files path to a tar ball of numpy
+        files, or loaded numpy data
     :param batch_size: the size of batches to create for the iterator
     :param iter_steps: the number of steps (batches) to create.
         Set to -1 for infinite, 0 for running through the loaded data once,

src/sparseml/onnx/utils/model.py

@@ -617,15 +617,16 @@ class DeepSparseModelRunner(_DeepSparseBaseModelRunner):
 
     :param model: the path to the ONNX model file or the loaded onnx.ModelProto
     :param batch_size: the size of the batch to create the model for
-    :param num_cores: the number of physical cores to run the model on
+    :param num_cores: the number of physical cores to run the model on. Defaults
+        to run on all available cores
     :param loss: the loss function, if any, to run for evaluation of the model
     """
 
     def __init__(
         self,
         model: Union[str, ModelProto],
         batch_size: int,
-        num_cores: int = -1,
+        num_cores: int = None,
         loss: Union[
             Callable[[Dict[str, numpy.ndarray], Dict[str, numpy.ndarray]], Any], None
         ] = None,
@@ -783,14 +784,15 @@ class DeepSparseAnalyzeModelRunner(_DeepSparseBaseModelRunner):
 
     :param model: the path to the ONNX model file or the loaded onnx.ModelProto
     :param batch_size: the size of the batch to create the model for
-    :param num_cores: the number of physical cores to run the model on
+    :param num_cores: the number of physical cores to run the model on. Defaults
+        to run on all available cores
     """
 
     def __init__(
         self,
         model: Union[str, ModelProto],
         batch_size: int,
-        num_cores: int = -1,
+        num_cores: int = None,
     ):
         super().__init__(model, batch_size, num_cores, loss=None)
 

src/sparseml/pytorch/models/classification/mnist.py

@@ -132,7 +132,7 @@ def forward(self, inp: Tensor):
     domain="cv",
     sub_domain="classification",
     architecture="mnistnet",
-    sub_architecture="none",
+    sub_architecture=None,
     default_dataset="mnist",
     default_desc="base",
 )

src/sparseml/utils/helpers.py

@@ -19,7 +19,6 @@
 
 import errno
 import fnmatch
-import glob
 import logging
 import os
 import re
@@ -30,6 +29,8 @@
 
 import numpy
 
+from sparsezoo.utils import load_numpy_list
+
 
 __all__ = [
     "ALL_TOKEN",
@@ -487,20 +488,22 @@ def load_labeled_data(
     Assumes sorted ordering for on disk. Will match between when a file glob is passed
     for either data and/or labels.
 
-    :param data: the file glob or list of arrays to use for data
-    :param labels: the file glob or list of arrays to use for labels, if any
+    :param data: the file glob, file path to numpy data tar ball, or list of arrays to
+        use for data
+    :param labels: the file glob, file path to numpy data tar ball, or list of arrays
+        to use for labels, if any
     :param raise_on_error: True to raise on any error that occurs;
         False to log a warning, ignore, and continue
     :return: a list containing tuples of the data, labels. If labels was passed in
         as None, will now contain a None for the second index in each tuple
     """
     if isinstance(data, str):
-        data = sorted(glob.glob(data))
+        data = load_numpy_list(data)
 
     if labels is None:
         labels = [None for _ in range(len(data))]
     elif isinstance(labels, str):
-        labels = sorted(glob.glob(labels))
+        labels = load_numpy_list(labels)
 
     if len(data) != len(labels) and labels:
         # always raise this error, lengths must match

tests/sparseml/onnx/helpers.py

@@ -753,7 +753,7 @@ def analyzer_models(request):
             {
                 "domain": "cv",
                 "sub_domain": "classification",
-                "architecture": "resnet-v1",
+                "architecture": "resnet_v1",
                 "sub_architecture": "50",
                 "framework": "pytorch",
                 "repo": "sparseml",
@@ -769,7 +769,7 @@ def analyzer_models(request):
             {
                 "domain": "cv",
                 "sub_domain": "classification",
-                "architecture": "mobilenet-v1",
+                "architecture": "mobilenet_v1",
                 "sub_architecture": "1.0",
                 "framework": "pytorch",
                 "repo": "sparseml",
@@ -786,14 +786,14 @@ def analyzer_models(request):
 def onnx_repo_models(request) -> OnnxRepoModelFixture:
     model_args, model_name = request.param
     model = Zoo.load_model(**model_args)
-    model_path = model.download_onnx_file(overwrite=False)
-    data_paths = model.download_data_files(overwrite=False)
+    model_path = model.onnx_file.downloaded_path()
+    data_paths = [data_file.downloaded_path() for data_file in model.data.values()]
 
     input_paths = None
     output_paths = None
     for path in data_paths:
-        if "_sample-inputs" in path:
-            input_paths = path.split(".tar")[0]
-        elif "_sample-outputs" in path:
-            output_paths = path.split(".tar")[0]
+        if "sample-inputs" in path:
+            input_paths = path
+        elif "sample-outputs" in path:
+            output_paths = path
     return OnnxRepoModelFixture(model_path, model_name, input_paths, output_paths)

tests/sparseml/onnx/optim/quantization/test_quantize_model_post_training.py

@@ -93,7 +93,7 @@ def test_quantize_model_post_training_mnist():
         sub_domain="classification",
         architecture="mnistnet",
         framework="pytorch",
-    )[0].download_onnx_file()
+    )[0].onnx_file.downloaded_path()
     quant_model_path = tempfile.NamedTemporaryFile(suffix=".onnx", delete=False).name
 
     # Prepare sample validation dataset
@@ -132,7 +132,7 @@ def test_quantize_model_post_training_resnet50_imagenette():
     resnet50_imagenette_path = Zoo.load_model(
         domain="cv",
         sub_domain="classification",
-        architecture="resnet-v1",
+        architecture="resnet_v1",
         sub_architecture="50",
         framework="pytorch",
         repo="sparseml",
@@ -141,7 +141,7 @@ def test_quantize_model_post_training_resnet50_imagenette():
         optim_name="base",
         optim_category="none",
         optim_target=None,
-    ).download_onnx_file(overwrite=True)
+    ).onnx_file.downloaded_path()
     quant_model_path = tempfile.NamedTemporaryFile(suffix=".onnx", delete=False).name
 
     # Prepare sample validation dataset

tests/sparseml/onnx/optim/test_analyzer_model.py

@@ -37,7 +37,7 @@
             {
                 "domain": "cv",
                 "sub_domain": "classification",
-                "architecture": "resnet-v1",
+                "architecture": "resnet_v1",
                 "sub_architecture": "50",
                 "framework": "pytorch",
                 "repo": "sparseml",
@@ -55,7 +55,7 @@ def analyzer_models_repo(request):
     model_args, output_path = request.param
     output_path = os.path.join(RELATIVE_PATH, "test_analyzer_model_data", output_path)
     model = Zoo.load_model(**model_args)
-    model_path = model.download_onnx_file(overwrite=False)
+    model_path = model.onnx_file.downloaded_path()
 
     if GENERATE_TEST_FILES:
         analyzer = ModelAnalyzer(model_path)

tests/sparseml/onnx/optim/test_analyzer_model_data/resnet50pytorch.json

@@ -175,6 +175,7 @@
       "bias_name": null,
       "bias_shape": null,
       "attributes": {
+        "ceil_mode": 0,
         "kernel_shape": [
           3,
           3
@@ -8014,35 +8015,12 @@
     },
     {
       "id": "494",
-      "op_type": "Constant",
-      "input_names": [],
-      "output_names": [
-        "494"
-      ],
-      "input_shapes": null,
-      "output_shapes": null,
-      "params": 0,
-      "prunable": false,
-      "prunable_params": -1,
-      "prunable_params_zeroed": 0,
-      "prunable_equation_sensitivity": null,
-      "flops": null,
-      "weight_name": null,
-      "weight_shape": null,
-      "bias_name": null,
-      "bias_shape": null,
-      "attributes": {
-        "value": null
-      }
-    },
-    {
-      "id": "495",
       "op_type": "Shape",
       "input_names": [
         "493"
       ],
       "output_names": [
-        "495"
+        "494"
       ],
       "input_shapes": [
         [
@@ -8070,41 +8048,11 @@
       "attributes": {}
     },
     {
-      "id": "496",
-      "op_type": "Gather",
-      "input_names": [
-        "495",
-        "494"
-      ],
-      "output_names": [
-        "496"
-      ],
-      "input_shapes": [
-        [
-          4
-        ]
-      ],
-      "output_shapes": null,
-      "params": 0,
-      "prunable": false,
-      "prunable_params": -1,
-      "prunable_params_zeroed": 0,
-      "prunable_equation_sensitivity": null,
-      "flops": null,
-      "weight_name": null,
-      "weight_shape": null,
-      "bias_name": null,
-      "bias_shape": null,
-      "attributes": {
-        "axis": 0
-      }
-    },
-    {
-      "id": "497",
+      "id": "495",
       "op_type": "Constant",
       "input_names": [],
       "output_names": [
-        "497"
+        "495"
       ],
       "input_shapes": null,
       "output_shapes": null,
@@ -8123,20 +8071,21 @@
       }
     },
     {
-      "id": "498",
-      "op_type": "Unsqueeze",
+      "id": "496",
+      "op_type": "Gather",
       "input_names": [
-        "496"
+        "494",
+        "495"
       ],
       "output_names": [
-        "498"
+        "496"
       ],
-      "input_shapes": null,
-      "output_shapes": [
+      "input_shapes": [
         [
-          1
+          4
         ]
       ],
+      "output_shapes": null,
       "params": 0,
       "prunable": false,
       "prunable_params": -1,
@@ -8148,19 +8097,17 @@
       "bias_name": null,
       "bias_shape": null,
       "attributes": {
-        "axes": [
-          0
-        ]
+        "axis": 0
       }
     },
     {
-      "id": "499",
+      "id": "498",
       "op_type": "Unsqueeze",
       "input_names": [
-        "497"
+        "496"
       ],
       "output_names": [
-        "499"
+        "498"
       ],
       "input_shapes": null,
       "output_shapes": [
@@ -8188,16 +8135,12 @@
       "id": "500",
       "op_type": "Concat",
       "input_names": [
-        "498",
-        "499"
+        "498"
       ],
       "output_names": [
         "500"
       ],
       "input_shapes": [
-        [
-          1
-        ],
         [
           1
         ]