From f373acd0e216694125d5dccbadc5445f17639ab1 Mon Sep 17 00:00:00 2001 From: Benjamin Date: Sat, 13 Feb 2021 15:16:25 -0500 Subject: [PATCH 1/2] load recipes into Managers from sparsezoo stubs --- src/sparseml/keras/optim/manager.py | 24 +++++++++++++++------ src/sparseml/pytorch/optim/manager.py | 22 +++++++++++++------ src/sparseml/tensorflow_v1/optim/manager.py | 24 +++++++++++++++------ src/sparseml/utils/helpers.py | 21 ++++++++++++++++-- tests/sparseml/utils/test_helpers.py | 18 ++++++++++++++++ 5 files changed, 87 insertions(+), 22 deletions(-) diff --git a/src/sparseml/keras/optim/manager.py b/src/sparseml/keras/optim/manager.py index ebb02013fb1..d529994314a 100644 --- a/src/sparseml/keras/optim/manager.py +++ b/src/sparseml/keras/optim/manager.py @@ -37,17 +37,27 @@ class ScheduledModifierManager(BaseManager, Modifier): """ @staticmethod - def from_yaml(file_path: str, add_modifiers: List[Modifier] = None): + def from_yaml( + file_path: str, + add_modifiers: List[Modifier] = None, + zoo_recipe_type: Union[str, None] = None, + ): """ - Convenience function used to create the manager of multiple modifiers - from a yaml file. + Convenience function used to create the manager of multiple modifiers from a + recipe file. - :param file_path: the path to the yaml file to load the modifier from + :param file_path: the path to the recipe file to load the modifier from, can + also be a SparseZoo model stub preceded by 'zoo:' to load a recipe for + a model stored in SparseZoo. i.e. '/path/to/local/recipe.yaml', + 'zoo:model/stub/path' :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 + :param zoo_recipe_type: optional recipe type to specify when loading a recipe + from a SparseZoo stub leave None when loading from a local file. + i.e. 'original', 'transfer' + :return: ScheduledModifierManager() created from the recipe file """ - yaml_str = load_recipe_yaml_str(file_path) + yaml_str = load_recipe_yaml_str(file_path, zoo_recipe_type) modifiers = Modifier.load_list(yaml_str) if add_modifiers: modifiers.extend(add_modifiers) diff --git a/src/sparseml/pytorch/optim/manager.py b/src/sparseml/pytorch/optim/manager.py index 46a76fc383e..0597b3f1362 100644 --- a/src/sparseml/pytorch/optim/manager.py +++ b/src/sparseml/pytorch/optim/manager.py @@ -56,17 +56,27 @@ class ScheduledModifierManager(BaseManager, Modifier): """ @staticmethod - def from_yaml(file_path: str, add_modifiers: List[Modifier] = None): + def from_yaml( + file_path: str, + add_modifiers: List[Modifier] = None, + zoo_recipe_type: Union[str, None] = None, + ): """ Convenience function used to create the manager of multiple modifiers from a - yaml file. + recipe file. - :param file_path: the path to the yaml file to load the modifier from + :param file_path: the path to the recipe file to load the modifier from, can + also be a SparseZoo model stub preceded by 'zoo:' to load a recipe for + a model stored in SparseZoo. i.e. '/path/to/local/recipe.yaml', + 'zoo:model/stub/path' :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 + :param zoo_recipe_type: optional recipe type to specify when loading a recipe + from a SparseZoo stub leave None when loading from a local file. + i.e. 'original', 'transfer' + :return: ScheduledModifierManager() created from the recipe file """ - yaml_str = load_recipe_yaml_str(file_path) + yaml_str = load_recipe_yaml_str(file_path, zoo_recipe_type) modifiers = Modifier.load_list(yaml_str) if add_modifiers: diff --git a/src/sparseml/tensorflow_v1/optim/manager.py b/src/sparseml/tensorflow_v1/optim/manager.py index b1098e5637d..106c44f6314 100644 --- a/src/sparseml/tensorflow_v1/optim/manager.py +++ b/src/sparseml/tensorflow_v1/optim/manager.py @@ -74,17 +74,27 @@ class ScheduledModifierManager(BaseManager, Modifier): """ @staticmethod - def from_yaml(file_path: str, add_modifiers: List[Modifier] = None): + def from_yaml( + file_path: str, + add_modifiers: List[Modifier] = None, + zoo_recipe_type: Union[str, None] = None, + ): """ - Convenience function used to create the manager of multiple modifiers - from a yaml file. + Convenience function used to create the manager of multiple modifiers from a + recipe file. - :param file_path: the path to the yaml file to load the modifier from + :param file_path: the path to the recipe file to load the modifier from, can + also be a SparseZoo model stub preceded by 'zoo:' to load a recipe for + a model stored in SparseZoo. i.e. '/path/to/local/recipe.yaml', + 'zoo:model/stub/path' :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 + :param zoo_recipe_type: optional recipe type to specify when loading a recipe + from a SparseZoo stub leave None when loading from a local file. + i.e. 'original', 'transfer' + :return: ScheduledModifierManager() created from the recipe file """ - yaml_str = load_recipe_yaml_str(file_path) + yaml_str = load_recipe_yaml_str(file_path, zoo_recipe_type) modifiers = Modifier.load_list(yaml_str) if add_modifiers: modifiers.extend(add_modifiers) diff --git a/src/sparseml/utils/helpers.py b/src/sparseml/utils/helpers.py index eaa9bf1bc6b..5ebe728364b 100644 --- a/src/sparseml/utils/helpers.py +++ b/src/sparseml/utils/helpers.py @@ -29,6 +29,7 @@ import numpy +from sparsezoo import Zoo from sparsezoo.utils import load_numpy_list @@ -763,16 +764,32 @@ def _tensors_export_batch( ) -def load_recipe_yaml_str(file_path: str) -> str: +def load_recipe_yaml_str( + file_path: str, zoo_recipe_type: Union[str, None] = None +) -> str: """ Loads a YAML recipe file to a string or extracts recipe from YAML front matter in a sparsezoo markdown recipe card. 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:'. i.e. + '/path/to/local/recipe.yaml', 'zoo:model/stub/path' + :param zoo_recipe_type: optional recipe type to specify when loading a recipe + from a SparseZoo stub leave None when loading from a local file. + i.e. 'original', 'transfer' :return: the recipe YAML configuration loaded as a string """ + if file_path.startswith("zoo:"): + zoo_kwargs = {"recipe_type": zoo_recipe_type} if zoo_recipe_type else {} + file_path = Zoo.download_recipe_from_stub(file_path, **zoo_kwargs) + elif zoo_recipe_type is not None: + raise ValueError( + f"zoo_recipe_type was set to {zoo_recipe_type} but a SparseZoo stub was " + f"not specified by path {file_path}. Stub paths should be prefixed by 'zoo:'" + ) 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..a201c7a0ea8 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,20 @@ 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,recipe_type", + [ + ( + "zoo:cv/classification/resnet_v1-50/pytorch/sparseml/imagenette/pruned-conservative", + None, + ), + ( + "zoo:cv/classification/resnet_v1-50/pytorch/sparseml/imagenette/pruned-conservative", + "original", + ), + ], +) +def test_load_recipe_yaml_str_zoo(zoo_path, recipe_type): + assert load_recipe_yaml_str(zoo_path, recipe_type) From df7ec1309fb1cf0f50c402b808d9e5bee4696461 Mon Sep 17 00:00:00 2001 From: Benjamin Date: Mon, 15 Feb 2021 21:46:59 -0500 Subject: [PATCH 2/2] moving recipe_type handling to SparseZoo only, supporting loading SparseZoo recipe objects --- src/sparseml/keras/optim/manager.py | 19 ++++++------- src/sparseml/pytorch/optim/manager.py | 19 ++++++------- src/sparseml/tensorflow_v1/optim/manager.py | 19 ++++++------- src/sparseml/utils/helpers.py | 31 ++++++++++----------- tests/sparseml/utils/test_helpers.py | 16 ++++------- 5 files changed, 47 insertions(+), 57 deletions(-) diff --git a/src/sparseml/keras/optim/manager.py b/src/sparseml/keras/optim/manager.py index d529994314a..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"] @@ -38,26 +39,24 @@ class ScheduledModifierManager(BaseManager, Modifier): @staticmethod def from_yaml( - file_path: str, + file_path: Union[str, OptimizationRecipe], add_modifiers: List[Modifier] = None, - zoo_recipe_type: Union[str, None] = None, ): """ 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, can - also be a SparseZoo model stub preceded by 'zoo:' to load a recipe for - a model stored in SparseZoo. i.e. '/path/to/local/recipe.yaml', - 'zoo:model/stub/path' + :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 recipe file - :param zoo_recipe_type: optional recipe type to specify when loading a recipe - from a SparseZoo stub leave None when loading from a local file. - i.e. 'original', 'transfer' :return: ScheduledModifierManager() created from the recipe file """ - yaml_str = load_recipe_yaml_str(file_path, zoo_recipe_type) + yaml_str = load_recipe_yaml_str(file_path) modifiers = Modifier.load_list(yaml_str) if add_modifiers: modifiers.extend(add_modifiers) diff --git a/src/sparseml/pytorch/optim/manager.py b/src/sparseml/pytorch/optim/manager.py index 0597b3f1362..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"] @@ -57,26 +58,24 @@ class ScheduledModifierManager(BaseManager, Modifier): @staticmethod def from_yaml( - file_path: str, + file_path: Union[str, OptimizationRecipe], add_modifiers: List[Modifier] = None, - zoo_recipe_type: Union[str, None] = None, ): """ 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, can - also be a SparseZoo model stub preceded by 'zoo:' to load a recipe for - a model stored in SparseZoo. i.e. '/path/to/local/recipe.yaml', - 'zoo:model/stub/path' + :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 recipe file - :param zoo_recipe_type: optional recipe type to specify when loading a recipe - from a SparseZoo stub leave None when loading from a local file. - i.e. 'original', 'transfer' :return: ScheduledModifierManager() created from the recipe file """ - yaml_str = load_recipe_yaml_str(file_path, zoo_recipe_type) + yaml_str = load_recipe_yaml_str(file_path) modifiers = Modifier.load_list(yaml_str) if add_modifiers: diff --git a/src/sparseml/tensorflow_v1/optim/manager.py b/src/sparseml/tensorflow_v1/optim/manager.py index 106c44f6314..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"] @@ -75,26 +76,24 @@ class ScheduledModifierManager(BaseManager, Modifier): @staticmethod def from_yaml( - file_path: str, + file_path: Union[str, OptimizationRecipe], add_modifiers: List[Modifier] = None, - zoo_recipe_type: Union[str, None] = None, ): """ 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, can - also be a SparseZoo model stub preceded by 'zoo:' to load a recipe for - a model stored in SparseZoo. i.e. '/path/to/local/recipe.yaml', - 'zoo:model/stub/path' + :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 recipe file - :param zoo_recipe_type: optional recipe type to specify when loading a recipe - from a SparseZoo stub leave None when loading from a local file. - i.e. 'original', 'transfer' :return: ScheduledModifierManager() created from the recipe file """ - yaml_str = load_recipe_yaml_str(file_path, zoo_recipe_type) + yaml_str = load_recipe_yaml_str(file_path) modifiers = Modifier.load_list(yaml_str) if add_modifiers: modifiers.extend(add_modifiers) diff --git a/src/sparseml/utils/helpers.py b/src/sparseml/utils/helpers.py index 5ebe728364b..8b41f30f422 100644 --- a/src/sparseml/utils/helpers.py +++ b/src/sparseml/utils/helpers.py @@ -30,6 +30,7 @@ import numpy from sparsezoo import Zoo +from sparsezoo.objects import OptimizationRecipe from sparsezoo.utils import load_numpy_list @@ -764,32 +765,30 @@ def _tensors_export_batch( ) -def load_recipe_yaml_str( - file_path: str, zoo_recipe_type: Union[str, None] = None -) -> 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 or stub to a SparseZoo model whose recipe will be downloaded and loaded. - SparseZoo stubs should be preceded by 'zoo:'. i.e. - '/path/to/local/recipe.yaml', 'zoo:model/stub/path' - :param zoo_recipe_type: optional recipe type to specify when loading a recipe - from a SparseZoo stub leave None when loading from a local file. - i.e. 'original', 'transfer' + 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 file_path.startswith("zoo:"): - zoo_kwargs = {"recipe_type": zoo_recipe_type} if zoo_recipe_type else {} - file_path = Zoo.download_recipe_from_stub(file_path, **zoo_kwargs) - elif zoo_recipe_type is not None: - raise ValueError( - f"zoo_recipe_type was set to {zoo_recipe_type} but a SparseZoo stub was " - f"not specified by path {file_path}. Stub paths should be prefixed by 'zoo:'" - ) + 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 a201c7a0ea8..e658a3c6345 100644 --- a/tests/sparseml/utils/test_helpers.py +++ b/tests/sparseml/utils/test_helpers.py @@ -102,17 +102,11 @@ def test_interpolate(x_cur, x0, x1, y0, y1, inter_func, out): @pytest.mark.parametrize( - "zoo_path,recipe_type", + "zoo_path", [ - ( - "zoo:cv/classification/resnet_v1-50/pytorch/sparseml/imagenette/pruned-conservative", - None, - ), - ( - "zoo:cv/classification/resnet_v1-50/pytorch/sparseml/imagenette/pruned-conservative", - "original", - ), + "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, recipe_type): - assert load_recipe_yaml_str(zoo_path, recipe_type) +def test_load_recipe_yaml_str_zoo(zoo_path): + assert load_recipe_yaml_str(zoo_path)