diff --git a/src/sparseml/keras/optim/manager.py b/src/sparseml/keras/optim/manager.py index ebb02013fb1..2b0a0998226 100644 --- a/src/sparseml/keras/optim/manager.py +++ b/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=' 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) diff --git a/src/sparseml/pytorch/optim/manager.py b/src/sparseml/pytorch/optim/manager.py index 46a76fc383e..ac6929d6ec4 100644 --- a/src/sparseml/pytorch/optim/manager.py +++ b/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=' 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) diff --git a/src/sparseml/tensorflow_v1/optim/manager.py b/src/sparseml/tensorflow_v1/optim/manager.py index b1098e5637d..5fefcc17615 100644 --- a/src/sparseml/tensorflow_v1/optim/manager.py +++ b/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=' 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) diff --git a/src/sparseml/utils/helpers.py b/src/sparseml/utils/helpers.py index eaa9bf1bc6b..8b41f30f422 100644 --- a/src/sparseml/utils/helpers.py +++ b/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=' 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( diff --git a/tests/sparseml/utils/test_helpers.py b/tests/sparseml/utils/test_helpers.py index 4d0d093d41b..e658a3c6345 100644 --- a/tests/sparseml/utils/test_helpers.py +++ b/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)