load recipes directly from SparseZoo

src/sparseml/keras/optim/manager.py

@@ -26,6 +26,7 @@
 from sparseml.keras.utils.logger import KerasLogger
 from sparseml.optim import BaseManager
 from sparseml.utils import load_recipe_yaml_str
+from sparsezoo.objects import OptimizationRecipe
 
 
 __all__ = ["ScheduledModifierManager"]
@@ -37,15 +38,23 @@ class ScheduledModifierManager(BaseManager, Modifier):
     """
 
     @staticmethod
-    def from_yaml(file_path: str, add_modifiers: List[Modifier] = None):
+    def from_yaml(
+        file_path: Union[str, OptimizationRecipe],
+        add_modifiers: List[Modifier] = None,
+    ):
         """
-        Convenience function used to create the manager of multiple modifiers
-        from a yaml file.
-
-        :param file_path: the path to the yaml file to load the modifier from
+        Convenience function used to create the manager of multiple modifiers from a
+        recipe file.
+
+        :param file_path: the path to the recipe file to load the modifier from, or
+            a SparseZoo model stub to load a recipe for a model stored in SparseZoo.
+            SparseZoo stubs should be preceded by 'zoo:', and can contain an optional
+            '?recipe_type=<type>' parameter. Can also be a SparseZoo OptimizationRecipe
+            object. i.e. '/path/to/local/recipe.yaml', 'zoo:model/stub/path',
+            'zoo:model/stub/path?recipe_type=transfer'
         :param add_modifiers: additional modifiers that should be added to the
-            returned manager alongside the ones loaded from the yaml file
-        :return: ScheduledModifierManager() created from the yaml file
+            returned manager alongside the ones loaded from the recipe file
+        :return: ScheduledModifierManager() created from the recipe file
         """
         yaml_str = load_recipe_yaml_str(file_path)
         modifiers = Modifier.load_list(yaml_str)

src/sparseml/pytorch/optim/manager.py

@@ -29,6 +29,7 @@
 from sparseml.pytorch.optim.modifier import Modifier, ScheduledModifier
 from sparseml.pytorch.utils import PyTorchLogger
 from sparseml.utils import load_recipe_yaml_str
+from sparsezoo.objects import OptimizationRecipe
 
 
 __all__ = ["ScheduledModifierManager", "load_manager"]
@@ -56,15 +57,23 @@ class ScheduledModifierManager(BaseManager, Modifier):
     """
 
     @staticmethod
-    def from_yaml(file_path: str, add_modifiers: List[Modifier] = None):
+    def from_yaml(
+        file_path: Union[str, OptimizationRecipe],
+        add_modifiers: List[Modifier] = None,
+    ):
         """
         Convenience function used to create the manager of multiple modifiers from a
-        yaml file.
-
-        :param file_path: the path to the yaml file to load the modifier from
+        recipe file.
+
+        :param file_path: the path to the recipe file to load the modifier from, or
+            a SparseZoo model stub to load a recipe for a model stored in SparseZoo.
+            SparseZoo stubs should be preceded by 'zoo:', and can contain an optional
+            '?recipe_type=<type>' parameter. Can also be a SparseZoo OptimizationRecipe
+            object. i.e. '/path/to/local/recipe.yaml', 'zoo:model/stub/path',
+            'zoo:model/stub/path?recipe_type=transfer'
         :param add_modifiers: additional modifiers that should be added to the
-            returned manager alongside the ones loaded from the yaml file
-        :return: ScheduledModifierManager() created from the yaml file
+            returned manager alongside the ones loaded from the recipe file
+        :return: ScheduledModifierManager() created from the recipe file
         """
         yaml_str = load_recipe_yaml_str(file_path)
         modifiers = Modifier.load_list(yaml_str)

src/sparseml/tensorflow_v1/optim/manager.py

@@ -25,6 +25,7 @@
 from sparseml.tensorflow_v1.optim.modifier import NM_RECAL, Modifier, ScheduledModifier
 from sparseml.tensorflow_v1.utils import tf_compat
 from sparseml.utils import load_recipe_yaml_str
+from sparsezoo.objects import OptimizationRecipe
 
 
 __all__ = ["ScheduledModifierManager"]
@@ -74,15 +75,23 @@ class ScheduledModifierManager(BaseManager, Modifier):
     """
 
     @staticmethod
-    def from_yaml(file_path: str, add_modifiers: List[Modifier] = None):
+    def from_yaml(
+        file_path: Union[str, OptimizationRecipe],
+        add_modifiers: List[Modifier] = None,
+    ):
         """
-        Convenience function used to create the manager of multiple modifiers
-        from a yaml file.
-
-        :param file_path: the path to the yaml file to load the modifier from
+        Convenience function used to create the manager of multiple modifiers from a
+        recipe file.
+
+        :param file_path: the path to the recipe file to load the modifier from, or
+            a SparseZoo model stub to load a recipe for a model stored in SparseZoo.
+            SparseZoo stubs should be preceded by 'zoo:', and can contain an optional
+            '?recipe_type=<type>' parameter. Can also be a SparseZoo OptimizationRecipe
+            object. i.e. '/path/to/local/recipe.yaml', 'zoo:model/stub/path',
+            'zoo:model/stub/path?recipe_type=transfer'
         :param add_modifiers: additional modifiers that should be added to the
-            returned manager alongside the ones loaded from the yaml file
-        :return: ScheduledModifierManager() created from the yaml file
+            returned manager alongside the ones loaded from the recipe file
+        :return: ScheduledModifierManager() created from the recipe file
         """
         yaml_str = load_recipe_yaml_str(file_path)
         modifiers = Modifier.load_list(yaml_str)

src/sparseml/utils/helpers.py

@@ -29,6 +29,8 @@
 
 import numpy
 
+from sparsezoo import Zoo
+from sparsezoo.objects import OptimizationRecipe
 from sparsezoo.utils import load_numpy_list
 
 
@@ -763,16 +765,30 @@ def _tensors_export_batch(
     )
 
 
-def load_recipe_yaml_str(file_path: str) -> str:
+def load_recipe_yaml_str(file_path: Union[str, OptimizationRecipe]) -> str:
     """
     Loads a YAML recipe file to a string or
     extracts recipe from YAML front matter in a sparsezoo markdown recipe card.
+    Recipes can also be provided as SparseZoo model stubs or OptimizationRecipe
+    objects.
 
     YAML front matter: https://jekyllrb.com/docs/front-matter/
 
-    :param file_path: file path to recipe YAML file or markdown recipe card
+    :param file_path: file path to recipe YAML file or markdown recipe card or
+        stub to a SparseZoo model whose recipe will be downloaded and loaded.
+        SparseZoo stubs should be preceded by 'zoo:', and can contain an optional
+        '?recipe_type=<type>' parameter. Can also be a SparseZoo OptimizationRecipe
+        object. i.e. '/path/to/local/recipe.yaml', 'zoo:model/stub/path',
+        'zoo:model/stub/path?recipe_type=transfer'
     :return: the recipe YAML configuration loaded as a string
     """
+    if isinstance(file_path, OptimizationRecipe):
+        # download and unwrap OptimizationRecipe object
+        file_path = file_path.downloaded_path()
+    elif file_path.startswith("zoo:"):
+        # download from zoo stub
+        file_path = Zoo.download_recipe_from_stub(file_path)
+
     extension = file_path.lower().split(".")[-1]
     if extension not in ["md", "yaml"]:
         raise ValueError(

tests/sparseml/utils/test_helpers.py

@@ -19,6 +19,7 @@
     convert_to_bool,
     flatten_iterable,
     interpolate,
+    load_recipe_yaml_str,
     validate_str_iterable,
 )
 
@@ -98,3 +99,14 @@ def test_validate_str_iterable_negative():
 def test_interpolate(x_cur, x0, x1, y0, y1, inter_func, out):
     interpolated = interpolate(x_cur, x0, x1, y0, y1, inter_func)
     assert abs(out - interpolated) < 0.01
+
+
+@pytest.mark.parametrize(
+    "zoo_path",
+    [
+        "zoo:cv/classification/resnet_v1-50/pytorch/sparseml/imagenette/pruned-conservative",
+        "zoo:cv/classification/resnet_v1-50/pytorch/sparseml/imagenette/pruned-conservative?recipe_type=original",
+    ],
+)
+def test_load_recipe_yaml_str_zoo(zoo_path):
+    assert load_recipe_yaml_str(zoo_path)