# Copyright (c) 2021 - present / Neuralmagic, Inc. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""
Classes for building YAML SparseML recipes without instantiating specific modifier\
implementations
"""
import textwrap
from copy import deepcopy
from typing import Any, Dict, List, Optional, Type, Union
import yaml
from sparseml.optim import BaseModifier, ModifierProp
from sparseml.sparsification.model_info import ModelInfo
from sparseml.sparsification.modifier_epoch import EpochRangeModifier
from sparseml.sparsification.modifier_lr import SetLearningRateModifier
from sparseml.sparsification.modifier_pruning import GMPruningModifier
from sparseml.utils import create_parent_dirs
__all__ = [
"ModifierYAMLBuilder",
"RecipeYAMLBuilder",
"PruningRecipeBuilder",
"to_yaml_str",
]
[docs]class ModifierYAMLBuilder(object):
"""
Class for building a YAML string representation of a modifier by setting
various properties of it. Properties are automatically inferred through the
serializable ModifierProps of the given modifier. They can be accessed through
auto-generated set_{name} and get_{name}.
:param modifier_class: reference to the class of modifier this object should create
a YAML representation for
:param kwargs: modifier property kwargs to values to initialize them. each key must
be a valid serializable ModifierProp of the modifier class
"""
def __init__(self, modifier_class: Type[BaseModifier], **kwargs):
assert issubclass(
modifier_class, BaseModifier
), "a subclass of Modifier must be used to instantiate a ModifierYAMLBuilder"
self._modifier_class = modifier_class
self._modifier_property_names = set()
self._properties = {}
for attr in dir(modifier_class):
attr_obj = getattr(modifier_class, attr)
if isinstance(attr_obj, ModifierProp) and attr_obj.serializable:
self._modifier_property_names.add(attr)
for key, value in kwargs.items():
if key in self._modifier_property_names:
self._properties[key] = value
else:
raise ValueError(
f"Modifier {modifier_class} has no serializable " f"property {key}"
)
def __getattr__(self, item: str) -> Any:
if item in self.__dict__:
return getattr(self, item)
elif item in self._modifier_property_names:
return self._properties.get(item)
elif item == "__name__":
return f"{self.__class__.__name__}.{self._modifier_class.__name__}"
else:
raise ValueError(
f"{self.__class__.__name__} of {self._modifier_class} has no "
f"property {item}"
)
def __setattr__(self, key: str, value: Any):
if key in ["_modifier_class", "_modifier_property_names", "_properties"]:
super().__setattr__(key, value)
elif key in self._modifier_property_names:
self._properties[key] = value
else:
raise ValueError(
f"{self.__class__.__name__} of {self._modifier_class.__name__} has no "
f"property {key}"
)
[docs] def copy(self) -> "ModifierYAMLBuilder":
"""
:return: newly constructed ModifierYAMLBuilder with the same base class and
properties
"""
properties = deepcopy(self._properties)
return self.__class__(self.modifier_class, **properties)
@property
def modifier_class(self) -> Type[BaseModifier]:
"""
:return: the class of the Modifier for which this object is building a string
"""
return self._modifier_class
[docs] def build_yaml_str(self) -> str:
"""
:return: string representation of the built Modifier as a YAML list item
"""
class_name_yaml = f"- !{self._modifier_class.__name__}"
properties_yaml = "\n".join(
[f"{key}: {to_yaml_str(value)}" for key, value in self._properties.items()]
)
properties_yaml = textwrap.indent(properties_yaml, " ")
return f"{class_name_yaml}\n{properties_yaml}"
[docs]class RecipeYAMLBuilder(object):
"""
Class for building a YAML SparseML recipe with standardized structure
:param variables: dict of string initial variable names to non-modifier recipe
variables to be included. Default is an empty dict
:param modifier_groups: dict of string initial modifier group names to a list
of ModifierYAMLBuilder objects of modifiers to be included in that group.
All modifier group names must contain 'modifiers' in the string. Default is
an empty dict
"""
def __init__(
self,
variables: Dict[str, Any] = None,
modifier_groups: Dict[str, List[ModifierYAMLBuilder]] = None,
):
self._variables = variables or {}
self._modifier_groups = modifier_groups or {}
self._validate()
[docs] def add_modifier_group(
self, name: str, modifier_builders: List[ModifierYAMLBuilder] = None
) -> "RecipeYAMLBuilder":
"""
Adds a modifier group with the given name to this builder
:param name: name of new modifier group
:param modifier_builders: list of modifier builder objects to initialize
this group with. Default is an empty list
:return: a reference to this object with the modifier group now added
"""
self._validate_modifier_group_name(name)
if name in self._modifier_groups:
raise KeyError(
f"{name} is already a modifier group name in this RecipeYAMLBuilder"
)
modifier_builders = modifier_builders or []
self._modifier_groups[name] = modifier_builders
self._validate()
return self
[docs] def get_modifier_group(self, name: str) -> Optional[List[ModifierYAMLBuilder]]:
"""
:param name: name of the modifier group to retrieve the modifier builders of
:return: reference to the list of modifier builders currently in this
modifier group. if the modifier group does not exist, None will be
returned
"""
return self._modifier_groups.get(name)
[docs] def get_modifier_builders(
self,
modifier_type: Optional[Union[Type[BaseModifier], str]] = None,
modifier_groups: Optional[Union[List[str], str]] = None,
):
"""
:param modifier_type: optional type of modifier to filter by. Can be
a type reference that will match if the modifier is of that type
or a subclass of it or a string where it will match if the class
is exactly that name. Defaults to None
:param modifier_groups: optional list of modifier group names to match
to. Defaults to None
:return: all modifier builders in this recipe, filtered by type and group
"""
if isinstance(modifier_groups, str):
modifier_groups = [modifier_groups]
modifier_builders = []
for group, builders in self._modifier_groups.items():
if modifier_groups is not None and group not in modifier_groups:
continue
for builder in builders:
if modifier_type and not self._modifier_builder_is_instance(
builder, modifier_type
):
continue
modifier_builders.append(builder)
return modifier_builders
[docs] def get_variable(self, name: str, default: Any = None) -> Any:
"""
:param name: name of the recipe variable to return
:param default: default value that should be returned if the given
name is not a current variable
:return: current value of the given variable, or the default if
the variable is not set in this builder
"""
return self._variables.get(name, default)
[docs] def has_variable(self, name: str) -> bool:
"""
:param name: name of the recipe variable to check
:return: True if this recipe builder has a variable with the given name.
False otherwise
"""
return name in self._variables
[docs] def set_variable(self, name: str, val: Any) -> "RecipeYAMLBuilder":
"""
Sets the given variable name to the given value
:param name: variable name to set
:param val: value to set the variable to
:return: a reference to this object with the variable now set
"""
self._variables[name] = val
return self
[docs] def build_yaml_str(self) -> str:
"""
:return: yaml string representation of this recipe in standard format
"""
# write variables
yaml_str = "\n".join(
[f"{key}: {to_yaml_str(value)}" for key, value in self._variables.items()]
)
# write modifier groups
for group, builders in self._modifier_groups.items():
if not builders:
continue # do not write empty groups
modifiers_yaml = "\n\n".join(
[builder.build_yaml_str() for builder in builders]
)
modifiers_yaml = textwrap.indent(modifiers_yaml, " ")
yaml_str += f"\n\n{group}:\n{modifiers_yaml}"
return yaml_str
[docs] def save_yaml(self, file_path: str):
"""
Saves this recipe as a yaml file to the given path
:param file_path: file path to save file to. if no '.' character is found
in the path, '.yaml' will be added to the path
"""
if "." not in file_path:
file_path += ".yaml"
self._save_file_str(self.build_yaml_str(), file_path)
[docs] def save_markdown(self, file_path: str, desc: str = ""):
"""
Saves this recipe as a markdown file to the given path with the
recipe yaml contained in the frontmatter
:param file_path: file path to save file to. if no '.' character is found
in the path, '.md' will be added to the path
:param desc: optional description to add to the markdown file after the recipe
YAML in the frontmatter. Default is empty string
"""
if "." not in file_path:
file_path += ".md"
md_content = f"---\n{self.build_yaml_str()}\n---\n{desc}"
self._save_file_str(md_content, file_path)
@staticmethod
def _save_file_str(content: str, file_path: str):
create_parent_dirs(file_path)
with open(file_path, "w") as file:
file.write(content)
@staticmethod
def _validate_modifier_group_name(name: str) -> bool:
if "modifiers" not in name:
raise ValueError(
"modifier groups must contain 'modifiers' in their name received "
f"group with name: {name}"
)
@staticmethod
def _modifier_builder_is_instance(
builder: ModifierYAMLBuilder, type_: Union[Type[BaseModifier], str]
) -> bool:
builder_class = builder.modifier_class
if isinstance(type_, str):
return builder_class.__name__ == type_
return builder_class is type_ or issubclass(builder_class, type_)
def _validate(self):
if not isinstance(self._variables, Dict):
raise ValueError(
"RecipeYAMLBuilder variables object must be a Dict "
f"found type {type(self._variables)}"
)
for name, builders in self._modifier_groups.items():
self._validate_modifier_group_name(name)
if not isinstance(builders, List):
raise ValueError(
"All modifier groups in RecipeYAMLBuilder must contain a list"
f"of ModifierYAMLBuilder objects. Group {name} has value of "
f"type {type(builders)}"
)
for builder in builders:
if not isinstance(builder, ModifierYAMLBuilder):
raise ValueError(
"All modifier groups in RecipeYAMLBuilder must contain a "
f"list of ModifierYAMLBuilder objects. Group {name} "
f"contains an element of type {type(builder)}"
)
[docs]class PruningRecipeBuilder(RecipeYAMLBuilder):
"""
Builds a basic, editable pruning recipe based on a given model info
standardized variables may be modified by constructor, or later on
| Sample yaml:
| num_epochs: 100
| init_lr: 0.0001
| pruning_start_target: 0.0
| pruning_end_target: 0.6
| pruning_update_frequency: 0.5
| base_target_sparsity: 0.8
| mask_type: unstructured
|
| training_modifiers:
| - !EpochRangeModifier
| start_epoch: 0.0
| end_epoch: eval(num_epochs)
|
| - !SetLearningRateModifier
| start_epoch: 0.0
| learning_rate: eval(init_lr)
|
| pruning_modifiers:
| - !GMPruningModifier
| params:
| - ... # based on prunable param names found in ModelInfo
| init_sparsity: 0.0
| final_sparsity: eval(base_target_sparsity)
| start_epoch: eval(pruning_start_target * num_epochs)
| end_epoch: eval(pruning_end_target * num_epochs)
| update_frequency: eval(pruning_update_frequency)
| mask_type: eval(mask_type)
:param model_info: model info object to extract layer information from
:param num_epochs: total number of epochs the recipe should run for. Default is 100
:param init_lr: initial learning rate value. Default is 0.0001
:param pruning_start_target: epoch that pruning should begin. this value
should be in range [0.0,1.0] representing the fraction of num_epochs
that the start epoch should be. (start_epoch=pruning_start_target*num_epochs).
Default is 0.0
:param pruning_end_target: epoch that pruning should complete. this value
should be in range [0.0,1.0] representing the fraction of num_epochs
that the end epoch should be. (end_epoch=pruning_end_target*num_epochs).
Default is 0.6
:param base_target_sparsity: target sparsity for pruning layers to. Default is 0.8
:param pruning_update_frequency: udpate frequency for pruning modifier.
Default is 0.5
:param mask_type: mask type to set the pruning modifier to. Default is unstructured
"""
def __init__(
self,
model_info: ModelInfo,
num_epochs: float = 100.0,
init_lr: float = 0.0001,
pruning_start_target: float = 0.0,
pruning_end_target: float = 0.6,
base_target_sparsity: float = 0.8,
pruning_update_frequency: float = 0.5,
mask_type: str = "unstructured",
):
self.num_epochs = num_epochs
self.init_lr = init_lr
self.pruning_start_target = pruning_start_target
self.pruning_end_target = pruning_end_target
self.pruning_update_frequency = pruning_update_frequency
self.base_target_sparsity = base_target_sparsity
self.mask_type = mask_type
super().__init__(
variables=dict(
num_epochs=self.num_epochs,
init_lr=self.init_lr,
pruning_start_target=self.pruning_start_target,
pruning_end_target=self.pruning_end_target,
pruning_update_frequency=self.pruning_update_frequency,
base_target_sparsity=self.base_target_sparsity,
mask_type=self.mask_type,
),
modifier_groups=dict(
training_modifiers=self._base_training_modifiers(),
pruning_modifiers=self._base_pruning_modifiers(model_info),
),
)
def __setattr__(self, key: str, value: Any):
# allow updates to base variables to propagate to the internal vars dict
if key in dir(self) and self.has_variable(key):
self.set_variable(key, value)
super().__setattr__(key, value)
@staticmethod
def _base_training_modifiers() -> List[ModifierYAMLBuilder]:
epoch_modifier = ModifierYAMLBuilder(
EpochRangeModifier, start_epoch=0.0, end_epoch="eval(num_epochs)"
)
init_lr_modifier = ModifierYAMLBuilder(
SetLearningRateModifier,
learning_rate="eval(init_lr)",
)
return [epoch_modifier, init_lr_modifier]
@staticmethod
def _base_pruning_modifiers(model_info: ModelInfo) -> List[ModifierYAMLBuilder]:
pruning_modifier = ModifierYAMLBuilder(
GMPruningModifier,
params=list(model_info.get_prunable_param_names()),
init_sparsity=0.0,
final_sparsity="eval(base_target_sparsity)",
start_epoch="eval(pruning_start_target * num_epochs)",
end_epoch="eval(pruning_end_target * num_epochs)",
update_frequency="eval(pruning_update_frequency)",
mask_type="eval(mask_type)",
)
return [pruning_modifier]
[docs] def build_yaml_str(self) -> str:
"""
:return: yaml string representation of this recipe in standard format
"""
for pruning_modifier in self.get_modifier_builders(GMPruningModifier):
params = pruning_modifier.params
if isinstance(params, list):
pruning_modifier.params = list(sorted(params))
return super().build_yaml_str()
[docs]def to_yaml_str(val: Any) -> str:
"""
:param val: value to get yaml str value of
:return: direct str cast of val if it is an int, float, or bool, otherwise
the stripped output of yaml.dump
"""
if isinstance(val, (str, int, float, bool)):
return str(val)
else:
yaml_str = yaml.dump(val).strip()
if isinstance(val, (Dict, List)):
yaml_str = "\n" + yaml_str
return yaml_str