SparseZoo V2 Bug Fixes (#1719)

* fixing naming, remove deprecated test * fix path to model.pth * debug * debug * fix for integration * debug * py3.9 compatability * Update helpers.py * updates for sparsezoo default change * update stub to one with multiple recipes
neuralmagic · Sep 15, 2023 · e5faef5 · e5faef5
1 parent cb1ccdd
commit e5faef5
Show file tree

Hide file tree

Showing 5 changed files with 14 additions and 17 deletions.
diff --git a/src/sparseml/pytorch/image_classification/utils/helpers.py b/src/sparseml/pytorch/image_classification/utils/helpers.py
@@ -91,9 +91,9 @@ def save_zoo_directory(
     """
     for root_file in [
         "model.onnx",
-        "sample_inputs",
-        "sample_outputs",
-        "sample_labels",
+        "sample-inputs",
+        "sample-outputs",
+        "sample-labels",
         "deployment",
     ]:
         root_file_path = os.path.join(training_outputs_dir, root_file)
@@ -109,9 +109,9 @@ def save_zoo_directory(
         training=os.path.join(training_outputs_dir, "training"),
         deployment=os.path.join(training_outputs_dir, "deployment"),
         onnx_model=os.path.join(training_outputs_dir, "model.onnx"),
-        sample_inputs=os.path.join(training_outputs_dir, "sample_inputs"),
-        sample_outputs=os.path.join(training_outputs_dir, "sample_outputs"),
-        sample_labels=os.path.join(training_outputs_dir, "sample_labels"),
+        sample_inputs=os.path.join(training_outputs_dir, "sample-inputs"),
+        sample_outputs=os.path.join(training_outputs_dir, "sample-outputs"),
+        sample_labels=os.path.join(training_outputs_dir, "sample-labels"),
         model_card=os.path.join(training_outputs_dir, "model.md"),
         logs=logs_path,
         sample_originals=None,

diff --git a/src/sparseml/pytorch/utils/exporter.py b/src/sparseml/pytorch/utils/exporter.py
@@ -362,9 +362,9 @@ def export_samples(
         :param exp_counter: the counter to start exporting the tensor files at
         """
         sample_batches = [tensors_to_device(batch, "cpu") for batch in sample_batches]
-        inputs_dir = os.path.join(self._output_dir, "sample_inputs")
-        outputs_dir = os.path.join(self._output_dir, "sample_outputs")
-        labels_dir = os.path.join(self._output_dir, "sample_labels")
+        inputs_dir = os.path.join(self._output_dir, "sample-inputs")
+        outputs_dir = os.path.join(self._output_dir, "sample-outputs")
+        labels_dir = os.path.join(self._output_dir, "sample-labels")
         originals_dir = os.path.join(self._output_dir, "sample_originals")
 
         with torch.no_grad():

diff --git a/src/sparseml/transformers/utils/helpers.py b/src/sparseml/transformers/utils/helpers.py
@@ -45,7 +45,7 @@ def save_zoo_directory(
         with the saved training artifacts
     :param logs_path: Optional directory where the training logs reside
     """
-    for root_file in ["sample_inputs", "sample_outputs"]:
+    for root_file in ["sample-inputs", "sample-outputs"]:
         root_file_path = os.path.join(training_outputs_dir, root_file)
         if not os.path.exists(root_file_path):
             logging.warning(
@@ -67,8 +67,8 @@ def save_zoo_directory(
         training=os.path.join(training_outputs_dir, "training"),
         deployment=os.path.join(training_outputs_dir, "deployment"),
         onnx_model=os.path.join(training_outputs_dir, "model.onnx"),
-        sample_inputs=os.path.join(training_outputs_dir, "sample_inputs"),
-        sample_outputs=os.path.join(training_outputs_dir, "sample_outputs"),
+        sample_inputs=os.path.join(training_outputs_dir, "sample-inputs"),
+        sample_outputs=os.path.join(training_outputs_dir, "sample-outputs"),
         model_card=os.path.join(training_outputs_dir, "model.md"),
         logs=logs_path,
         sample_labels=None,

diff --git a/src/sparseml/yolov8/utils/export_samples.py b/src/sparseml/yolov8/utils/export_samples.py
@@ -78,7 +78,7 @@ def export_sample_inputs_outputs(
     exported_samples = 0
 
     # Sample export directories
-    sample_in_dir = os.path.join(save_dir, "sample_inputs")
+    sample_in_dir = os.path.join(save_dir, "sample-inputs")
     sample_out_dir_torch = os.path.join(save_dir, "sample_outputs_torch")
     sample_out_dir_ort = os.path.join(save_dir, "sample_outputs_onnxruntime")
 

diff --git a/tests/sparseml/optim/test_helpers.py b/tests/sparseml/optim/test_helpers.py
@@ -627,10 +627,7 @@ def test_evaluate_recipe_yaml_str_equations_invalid(recipe):
             "zoo:cv/classification/resnet_v1-50/pytorch/sparseml/imagenette/"
             "pruned-conservative"
         ),
-        (
-            "zoo:cv/classification/resnet_v1-50/pytorch/sparseml/imagenette/"
-            "pruned-conservative?recipe=original"
-        ),
+        ("zoo:bert-base-wikipedia_bookcorpus-pruned90?recipe=transfer_question"),
     ],
 )
 def test_load_recipe_yaml_str_zoo(zoo_path):