sparseml.sparsification package
Submodules
sparseml.sparsification.analyzer module
Code for running analysis on neural networks
-
class
sparseml.sparsification.analyzer.
Analyzer
(model_info: sparseml.sparsification.model_info.ModelInfo)[source] Bases:
abc.ABC
Base abstract class for model analyzers. Analyzers should be able to detect if given a ModelInfo object and other keyword inputs if they should run their analysis.
- Parameters
model_info – ModelInfo object of the model to be analyzed. after running this analysis, the analysis_results of this ModelInfo object will be updated
-
abstract classmethod
available
(model_info: sparseml.sparsification.model_info.ModelInfo, **kwargs) → bool[source] Abstract method that subclasses must implement to determine if given the model info and keyword arguments that the Analyzer can run its analysis
- Parameters
model_info – ModelInfo object of the model to be analyzed
kwargs – additional keyword arguments that will be passed to the run function
- Returns
True if given the inputs, this analyzer can run its analysis. False otherwise
-
run
(show_progress: bool = False, **kwargs) → sparseml.sparsification.model_info.ModelResult[source] Runs the given analysis by calling to the underlying run_iter method :param show_progress: set True to display a tqdm progress bar. default is False :param kwargs: key word arguments validated by available() to run this analysis :return: the final result from this analysis. this result will also be
added to the ModelInfo object of this Analyzer
-
run_iter
(**kwargs) → Generator[Tuple[sparseml.sparsification.analyzer.AnalyzerProgress, sparseml.sparsification.model_info.ModelResult], None, None][source] runs the analysis stepwise using the abstract _run_iter method yielding an AnalyzerProgress and the in progress ModelResult at each step
After the last step, the final results will be added to the given ModelInfo
- Parameters
kwargs – key word arguments validated by available() to run this analysis
-
class
sparseml.sparsification.analyzer.
AnalyzerProgress
(*, step: int, total_steps: int, metadata: Dict[str, Any] = None)[source] Bases:
pydantic.main.BaseModel
Simple class for tracking model analyzer progress
-
metadata
: Optional[Dict[str, Any]]
-
property
progress
float progress on [0,1] scale
- Type
return
-
step
: int
-
total_steps
: int
-
-
class
sparseml.sparsification.analyzer.
PruningLossSensitivityMagnitudeAnalyzer
(model_info: sparseml.sparsification.model_info.ModelInfo)[source] Bases:
sparseml.sparsification.analyzer.Analyzer
,abc.ABC
Base class for running pruning loss sensitivity weight magnitude analysis. A valid in-framework model with prunable weights is required to run this analysis
pruning_loss_analysis_sparsity_levels is an optional run argument to set the sparsities that this analysis will run at. if not set, the value defaults to sparsml.optim.default_pruning_sparsities_loss(extended=True)
-
classmethod
available
(model_info: sparseml.sparsification.model_info.ModelInfo, **kwargs) → bool[source] Determines if given the available kwargs and ModelInfo, that weight magnitude analysis is available. model must exist in the given keyword arguments and be a valid model of the given framework and include all prunable parameters named in the ModelInfo
- Parameters
model_info – ModelInfo object of the model to be analyzed
kwargs – keyword arguments that will be passed in to this analysis. model must be included for this analysis to be available
- Returns
True if given the inputs, this analyzer can run its analysis. False otherwise
-
abstract
get_named_prunable_params
(model: Any) → Dict[str, numpy.ndarray][source] loads the prunable parameters in a standardized way so that weight magnitude analysis may be run on each
- Parameters
model – model to load the prunable parameters from
- Returns
dictionary of prunable parameter name as listed in the ModelInfo to a numpy array of the values of the parameter
-
abstract static
validate_model
(prunable_param_names: Set[str], model: Any) → bool[source] Validates that all prunable parameter names in the ModelInfo layer_info exist in the given model and that the given model is of the correct framework
- Parameters
prunable_param_names – set of prunable parameter names found in the model info
model – model to validate
- Returns
True if this is a valid model for weight mangitude pruning analysis. False otherwise
-
classmethod
sparseml.sparsification.info module
Functionality related to describing availability and information of sparsification algorithms to models within in the ML frameworks.
The file is executable and will get the sparsification info for a given framework:
Compile the available setup and information for the sparsification of a model in a given framework.
- positional arguments:
- framework the ML framework or path to a framework file to load the
sparsification info for
- optional arguments:
- -h, --help
show this help message and exit
- --path PATH
A full file path to save the sparsification info to. If not supplied, will print out the sparsification info to the console.
-
class
sparseml.sparsification.info.
ModifierInfo
(*, name: str, description: str, type_: sparseml.sparsification.info.ModifierType = <ModifierType.misc: 'misc'>, props: List[sparseml.sparsification.info.ModifierPropInfo] = [], warnings: List[str] = None)[source] Bases:
pydantic.main.BaseModel
Class for storing information and associated metadata for a given Modifier. Extends pydantics BaseModel class for serialization to and from json in addition to proper type checking on construction.
-
description
: str
-
name
: str
-
props
: List[sparseml.sparsification.info.ModifierPropInfo]
-
warnings
: Optional[List[str]]
-
-
class
sparseml.sparsification.info.
ModifierPropInfo
(*, name: str, description: str, type_: str, restrictions: List[Any] = None)[source] Bases:
pydantic.main.BaseModel
Class for storing information and associated metadata for a property on a given Modifier. Extends pydantics BaseModel class for serialization to and from json in addition to proper type checking on construction.
-
description
: str
-
name
: str
-
restrictions
: Optional[List[Any]]
-
type_
: str
-
-
class
sparseml.sparsification.info.
ModifierType
(value)[source] Bases:
enum.Enum
Types of modifiers for grouping what functionality a Modifier falls under.
-
act_sparsity
= 'act_sparsity'
-
general
= 'general'
-
misc
= 'misc'
-
pruning
= 'pruning'
-
quantization
= 'quantization'
-
training
= 'training'
-
-
class
sparseml.sparsification.info.
SparsificationInfo
(*, modifiers: List[sparseml.sparsification.info.ModifierInfo] = [])[source] Bases:
pydantic.main.BaseModel
Class for storing the information for sparsifying in a given framework. Extends pydantics BaseModel class for serialization to and from json in addition to proper type checking on construction.
-
modifiers
: List[sparseml.sparsification.info.ModifierInfo]
-
type_modifiers
(type_: sparseml.sparsification.info.ModifierType) → List[sparseml.sparsification.info.ModifierInfo][source] Get the contained Modifiers for a specific ModifierType. :param type_: The ModifierType to filter the returned list of Modifiers by. :type type_: ModifierType :return: The filtered list of Modifiers that match the given type_. :rtype: List[ModifierInfo]
-
-
sparseml.sparsification.info.
load_sparsification_info
(load: str) → sparseml.sparsification.info.SparsificationInfo[source] Load the sparsification info from a file or raw json. If load exists as a path, will read from the file and use that. Otherwise will try to parse the input as a raw json str.
- Parameters
load (str) – Either a file path to a json file or a raw json string.
- Returns
The loaded sparsification info.
- Return type
-
sparseml.sparsification.info.
save_sparsification_info
(framework: Any, path: Optional[str] = None)[source] Save the sparsification info for a given framework. If path is provided, will save to a json file at that path. If path is not provided, will print out the info.
- Parameters
framework (Any) – The item to detect the ML framework for. See
detect_framework()
for more information.path (Optional[str]) – The path, if any, to save the info to in json format. If not provided will print out the info.
-
sparseml.sparsification.info.
sparsification_info
(framework: Any) → sparseml.sparsification.info.SparsificationInfo[source] Get the available setup for sparsifying model in the given framework.
- Parameters
framework (Any) – The item to detect the ML framework for. See
detect_framework()
for more information.- Returns
The sparsification info for the given framework
- Return type
sparseml.sparsification.model_info module
Base classes for describing models and layers in ML framework neural networks.
-
class
sparseml.sparsification.model_info.
LayerInfo
(*, name: str, op_type: str, params: int = None, bias_params: int = None, prunable: bool = False, flops: int = None, execution_order: int = - 1, attributes: Dict[str, Any] = None)[source] Bases:
pydantic.main.BaseModel
Class for storing properties about a layer in a model
-
attributes
: Optional[Dict[str, Any]]
-
bias_params
: Optional[int]
-
classmethod
conv_layer
(name: str, in_channels: int, out_channels: int, kernel_shape: List[int], bias: bool, groups: int = 1, stride: Union[int, List[int]] = 1, padding: Optional[List[int]] = None, **kwargs) → sparseml.sparsification.model_info.LayerInfo[source] creates a LayerInfo object for a fully connected convolutional layer
- Parameters
name – unique name of the layer within its model
in_channels – number of input channels
out_channels – number of output channels
kernel_shape – kernel shape of this layer, given as a list
bias – True if the linear layer has a bias add included, False otherwise
groups – number of groups that input and output channels are divided into. default is 1
stride – stride for this convolution, can be int or tuple of ints. default is 1
padding – padding applied to each spatial axis. defualt is [0, 0, 0, 0]
kwargs – additional kwargs to be passed to the LayerInfo constructor
- Returns
-
execution_order
: int
-
flops
: Optional[int]
-
classmethod
linear_layer
(name: str, in_channels: int, out_channels: int, bias: bool, **kwargs) → sparseml.sparsification.model_info.LayerInfo[source] creates a LayerInfo object for a fully connected linear layer
- Parameters
name – unique name of the layer within its model
in_channels – number of input channels
out_channels – number of output channels
bias – True if the linear layer has a bias add included, False otherwise
kwargs – additional kwargs to be passed to the LayerInfo constructor
- Returns
-
name
: str
-
op_type
: str
-
params
: Optional[int]
-
prunable
: bool
-
-
class
sparseml.sparsification.model_info.
ModelInfo
(model: Any, metadata: Optional[Dict[str, Any]] = None)[source] Bases:
abc.ABC
Base class for extracting and serializing model metadata, layers info, and analysis results
- Parameters
model – framework specific model object to extract info for
metadata – optional dict of string metadata attributes to value. Default is empty dict
-
add_analysis_result
(result: sparseml.sparsification.model_info.ModelResult)[source]
-
property
analysis_results
list of analysis results run on this model
- Type
return
-
abstract
extract_layer_info
(model: Any) → OrderedDict[str, LayerInfo][source] Abstract method for extracting an ordered dictionary of layer name to completed LayerInfo object for the layer
- Parameters
model – model to extract LayerInfo information of
- Returns
ordered dictionary of layer name to LayerInfo object for the layer
-
classmethod
from_dict
(dictionary: Dict[str, Any])[source] - Parameters
dictionary – dict serialized by dict(ModelInfo(…))
- Returns
ModelInfo object created from the given dict
-
get_prunable_param_names
() → Set[str][source] - Returns
set of parameter names of all prunable layers in this ModelInfo
-
get_results_by_type
(analysis_type: str) → List[sparseml.sparsification.model_info.ModelResult][source] - Parameters
analysis_type – type of analysis in ModelResult.analysis_type to filter by
- Returns
list of analysis results of this model that match the given type
-
property
layer_info
dict of unique layer name to LayerInfo object of the given layer
- Type
return
-
static
load
(file_path) → sparseml.sparsification.model_info.ModelInfo[source] - Parameters
file_path – file path to JSON file to load ModelInfo object from
- Returns
the loaded ModelInfo object
-
class
sparseml.sparsification.model_info.
ModelResult
(*, value: Any = None, attributes: Dict[str, Any] = None, analysis_type: str, layer_results: Dict[str, sparseml.sparsification.model_info.Result] = None)[source] Bases:
sparseml.sparsification.model_info.Result
Class for storing the results of an analysis for an entire model
-
analysis_type
: str
-
layer_results
: Dict[str, sparseml.sparsification.model_info.Result]
-
-
class
sparseml.sparsification.model_info.
PruningSensitivityResult
(analysis_type: sparseml.sparsification.model_info.PruningSensitivityResultTypes, *, value: Any = None, attributes: Dict[str, Any] = None, layer_results: Dict[str, sparseml.sparsification.model_info.Result] = None)[source] Bases:
sparseml.sparsification.model_info.ModelResult
Helper class for creating and updating results of pruning sensitivity analyses
- Parameters
analysis_type – PruningSensitivityResultTypes Enum value of type of analysis this is
kwargs – optional args to be passed into model result constructor
-
add_layer_sparsity_result
(layer_name: str, sparsity: float, value: Any)[source] Adds a result of the given value for a given sparsity to the given layer name :param layer_name: layer param name to add result for :param sparsity: sparsity of the layer at which the sensitivity was measured :param value: sensitivity value
-
add_model_sparsity_result
(sparsity: float, value: Any)[source] Adds a model result of the given value for a given sparsity :param sparsity: sparsity of model at which the sensitivity was measured :param value: sensitivity value
-
analysis_type
: str
-
attributes
: Optional[Dict[str, Any]]
-
get_available_layer_sparsities
() → List[float][source] - Returns
list of sparsity values available for all model layers
-
get_layer_sparsity_score
(layer_name: str, sparsity: float) → float[source] - Parameters
layer_name – name of layer to get sparsity score for
sparsity – sparsity to measure score at
- Returns
sparsity score at the given sparsity such that higher scores correlate to a less prunable layer
-
layer_results
: Dict[str, sparseml.sparsification.model_info.Result]
-
value
: Any
sparseml.sparsification.modifier_epoch module
Base class for modifiers related to controlling the training epochs while training a model
-
class
sparseml.sparsification.modifier_epoch.
EpochRangeModifier
(start_epoch: float, end_epoch: float, **kwargs)[source] Bases:
sparseml.optim.modifier.BaseModifier
,sparseml.optim.modifier.BaseScheduled
Simple modifier to set the range of epochs when applying a modifier (ie to set min and max epochs within a range without hacking other modifiers).
Note, that if other modifiers exceed the range of this one for min or max epochs, this modifier will not have an effect.
Sample yaml:!EpochRangeModifier:start_epoch: 0end_epoch: 90- Parameters
start_epoch – The epoch to start the modifier at
end_epoch – The epoch to end the modifier at
-
sparsification_types
the sparsification types this modifier instance will apply
- Type
return
sparseml.sparsification.modifier_lr module
Code related to learning rate controls that are shared across frameworks.
-
class
sparseml.sparsification.modifier_lr.
LearningRateModifier
(lr_class: str, lr_kwargs: Dict, init_lr: float, start_epoch: float, end_epoch: float = - 1.0, update_frequency: float = - 1.0, **kwargs)[source] Bases:
sparseml.optim.modifier.BaseModifier
,sparseml.optim.modifier.BaseScheduled
,sparseml.optim.modifier.BaseUpdate
Generic implementation for LearningRateModifier shared across framework implementations.
Sample yaml:!LearningRateModifierlr_class: ExponentialDecaylr_kwargs:initial_learning_rate: 0.01decay_steps: 10000decay_rate: 0.96start_epoch: 0.0end_epoch: 10.0- Parameters
lr_class – The name of the lr scheduler class to use: [StepLR, MultiStepLR, ExponentialLR]
lr_kwargs – The dictionary of keyword arguments to pass to the constructor for the lr_class
init_lr – The initial learning rate to use once this modifier starts
start_epoch – The epoch to start the modifier at (set to -1.0 so it starts immediately)
end_epoch – The epoch to end the modifier at, (set to -1.0 so it doesn’t end)
update_frequency – unused and should not be set
-
corrected_lr_info
(steps_per_epoch: int, start_epoch: float, end_epoch: float) → Tuple[str, Dict][source] Get the corrected learning rate info for use with modifiers. Normalizes any epoch values to steps.
- Parameters
steps_per_epoch – number of steps taken within each epoch
start_epoch – The epoch the LR should start being controlled at
end_epoch – The epoch the LR should stop being controlled at
- Returns
a tuple containing the corrected lr class and keyword args
-
init_lr
The initial learning rate to use once this modifier starts
- Type
return
-
lr_class
The name of the lr scheduler class to use: [StepLR, MultiStepLR, ExponentialLR]
- Type
return
-
lr_kwargs
The dictionary of keyword arguments to pass to the constructor for the lr_class
- Type
return
-
sparsification_types
the sparsification types this modifier instance will apply
- Type
return
-
class
sparseml.sparsification.modifier_lr.
SetLearningRateModifier
(learning_rate: float, start_epoch: float = - 1.0, end_epoch: float = - 1.0, **kwargs)[source] Bases:
sparseml.optim.modifier.BaseModifier
,sparseml.optim.modifier.BaseScheduled
Generic implementation for SetLearningRateModifier shared across framework implementations.
Sample yaml:!SetLearningRateModifierstart_epoch: 0.0learning_rate: 0.001- Parameters
learning_rate – The learning rate to use once this modifier starts
start_epoch – The epoch to start the modifier at
- Param
end_epoch: should not be set, does not affect modifier. Set at -1
-
learning_rate
The learning rate to use once this modifier starts
- Type
return
-
sparsification_types
the sparsification types this modifier instance will apply
- Type
return
sparseml.sparsification.modifier_params module
Base Modifier for changing the state of a modules params while training according to certain update formulas or patterns.
-
class
sparseml.sparsification.modifier_params.
TrainableParamsModifier
(params: Union[str, List[str]], trainable: bool, params_strict: bool = True, start_epoch: float = - 1.0, end_epoch: float = - 1.0, **kwargs)[source] Bases:
sparseml.optim.modifier.BaseModifier
,sparseml.optim.modifier.BaseScheduled
Base Modifier to control the trainability for a given list of parameters.
To select all params in the graph, set to the ALL_TOKEN string: __ALL__
Sample yaml:!TrainableParamsModifier:params: [“conv_net/conv1/weight”]trainable: True- Parameters
params – A list of full parameter names or regex patterns of names to apply pruning to. Regex patterns must be specified with the prefix ‘re:’. __ALL__ will match to all parameters. Can also use the token __ALL__ to specify all params
trainable – True if the param(s) should be made trainable, False to make them non-trainable
params_strict – True if the given param(s) must be found in each layer and will raise an err if not found, False if missing params are ok and will not raise an err
start_epoch – The epoch to start the modifier at (set to -1.0 so it starts immediately)
end_epoch – The epoch to end the modifier at (set to -1.0 so it never ends), if > 0 then will revert to the original value for the params after this epoch
-
params
A list of full parameter names or regex patterns of names to apply pruning to. Regex patterns must be specified with the prefix ‘re:’. __ALL__ will match to all parameters. Can also use the token __ALL__ to specify all params
- Type
return
-
params_strict
True if the given param(s) must be found in each layer and will raise an err if not found, False if missing params are ok and will not raise an err
- Type
return
-
sparsification_types
the sparsification types this modifier instance will apply
- Type
return
-
trainable
True if the param(s) should be made trainable, False to make them non-trainable
- Type
return
sparseml.sparsification.modifier_pruning module
Base Modifiers for inducing / enforcing kernel sparsity (model pruning) on models while pruning.
-
class
sparseml.sparsification.modifier_pruning.
ConstantPruningModifier
(params: Union[str, List[str]], start_epoch: float = - 1, end_epoch: float = - 1, **kwargs)[source] Bases:
sparseml.optim.modifier.BaseModifier
,sparseml.optim.modifier.BaseScheduled
Base modifier for holding the sparsity level and shape for a given param constant while training. Useful for transfer learning use cases.
Sample yaml:!ConstantPruningModifierparams: __ALL__start_epoch: 0.0end_epoch: 10.0- Parameters
params – List of str names or regex patterns of names for the parameter variables to apply the KS modifier to. Regex patterns must be specified with the prefix ‘re:’. Can also use the token __ALL__ to specify all prunable layers and weights
start_epoch – The epoch to start the modifier at
end_epoch – The epoch to end the modifier at
-
params
List of str for the variable names or regex patterns of names to apply the pruning modifier to. Regex patterns must be specified with the prefix ‘re:’.
- Type
return
-
sparsification_types
the sparsification types this modifier instance will apply
- Type
return
-
class
sparseml.sparsification.modifier_pruning.
GMPruningModifier
(params: Union[str, List[str]], init_sparsity: float, final_sparsity: float, start_epoch: float, end_epoch: float, update_frequency: float, inter_func: str = 'cubic', mask_type: Union[str, List[int]] = 'unstructured', leave_enabled: bool = True, **kwargs)[source] Bases:
sparseml.optim.modifier.BaseModifier
,sparseml.optim.modifier.BaseScheduled
,sparseml.optim.modifier.BaseUpdate
Base Modifier to gradually applies kernel sparsity to a given variable or variables from init_sparsity until final_sparsity is reached over a given amount of time and applied with an interpolated function for each step taken.
Applies based on magnitude pruning without any structure to the pruning.
Sample yaml:!GMPruningModifierparams: __ALL__init_sparsity: 0.05final_sparsity: 0.8start_epoch: 0.0end_epoch: 10.0update_frequency: 1.0inter_func: cubicmask_type: unstructuredleave_enabled: True- Parameters
params – List of str names or name regex patterns for the variables in the graph to apply the KS modifier to. Regex patterns must be specified with the prefix ‘re:’. __ALL__ will match to all parameters.
init_sparsity – The initial sparsity for the variable to start with at start_epoch
final_sparsity – The final sparsity for the variable to end with at end_epoch
start_epoch – The epoch to start the modifier at
end_epoch – The epoch to end the modifier at
update_frequency – The number of epochs or fraction of epochs to update at between start and end
leave_enabled – True to continue masking the weights after end_epoch, False to stop masking. Should be set to False if exporting the result immediately after or doing some other prune
inter_func – The type of interpolation function to use: [linear, cubic, inverse_cubic]
mask_type – String to define type of sparsity (options: [‘unstructured’, ‘channel’, ‘filter’]), List to define block shape of a parameter’s in and out channels. default is ‘unstructured’
leave_enabled – True to continue masking the weights after end_epoch, False to stop masking. Should be set to False if exporting the result immediately after or doing some other prune
-
final_sparsity
The final sparsity for the variable to end with at end_epoch
- Type
return
-
init_sparsity
The initial sparsity for the variable to start with at start_epoch
- Type
return
-
inter_func
The type of interpolation function to use: [linear, cubic, inverse_cubic]
- Type
return
-
leave_enabled
True to continue masking the weights after end_epoch, False to stop masking. Note, if set as False, sparsity will not be enforced and the model will likely deviate from the sparse solution
- Type
return
-
mask_type
the mask type used
- Type
return
-
params
List of str for the variable names or regex patterns of names to apply the KS modifier to. Regex patterns must be specified with the prefix ‘re:’.
- Type
return
-
sparsification_types
the sparsification types this modifier instance will apply
- Type
return
sparseml.sparsification.oracle module
Sparsification oracle functions for generating SparseML recipes from models
-
sparseml.sparsification.oracle.
create_pruning_recipe
(model: Any, save_path: Optional[str] = None, analyzer_kwargs: Optional[Dict[str, Any]] = None, skip_analyzer_types: Optional[List[Type[sparseml.sparsification.analyzer.Analyzer]]] = None) → Optional[str][source] - Parameters
model – loaded framework model or model file path of a model to create a recipe for
save_path – optional path to save the created recipe to
analyzer_kwargs – keyword arguments to be passed to the available() and run() functions of analyzer objects
skip_analyzer_types – list of Analyzer class types not to run even if available
- Returns
string of the created recipe if None is provided
sparseml.sparsification.recipe_builder module
Classes for building YAML SparseML recipes without instantiating specific modifierimplementations
-
class
sparseml.sparsification.recipe_builder.
ModifierYAMLBuilder
(modifier_class: Type[sparseml.optim.modifier.BaseModifier], **kwargs)[source] Bases:
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}.
- Parameters
modifier_class – reference to the class of modifier this object should create a YAML representation for
kwargs – modifier property kwargs to values to initialize them. each key must be a valid serializable ModifierProp of the modifier class
-
build_yaml_str
() → str[source] - Returns
string representation of the built Modifier as a YAML list item
-
copy
() → sparseml.sparsification.recipe_builder.ModifierYAMLBuilder[source] - Returns
newly constructed ModifierYAMLBuilder with the same base class and properties
-
property
modifier_class
the class of the Modifier for which this object is building a string
- Type
return
-
class
sparseml.sparsification.recipe_builder.
PruningRecipeBuilder
(model_info: sparseml.sparsification.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')[source] Bases:
sparseml.sparsification.recipe_builder.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: 100init_lr: 0.0001pruning_start_target: 0.0pruning_end_target: 0.6pruning_update_frequency: 0.5base_target_sparsity: 0.8mask_type: unstructuredtraining_modifiers:- !EpochRangeModifierstart_epoch: 0.0end_epoch: eval(num_epochs)- !SetLearningRateModifierstart_epoch: 0.0learning_rate: eval(init_lr)pruning_modifiers:- !GMPruningModifierparams:- … # based on prunable param names found in ModelInfoinit_sparsity: 0.0final_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)- Parameters
model_info – model info object to extract layer information from
num_epochs – total number of epochs the recipe should run for. Default is 100
init_lr – initial learning rate value. Default is 0.0001
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
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
base_target_sparsity – target sparsity for pruning layers to. Default is 0.8
pruning_update_frequency – udpate frequency for pruning modifier. Default is 0.5
mask_type – mask type to set the pruning modifier to. Default is unstructured
-
class
sparseml.sparsification.recipe_builder.
RecipeYAMLBuilder
(variables: Optional[Dict[str, Any]] = None, modifier_groups: Optional[Dict[str, List[sparseml.sparsification.recipe_builder.ModifierYAMLBuilder]]] = None)[source] Bases:
object
Class for building a YAML SparseML recipe with standardized structure
- Parameters
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
-
add_modifier_group
(name: str, modifier_builders: Optional[List[sparseml.sparsification.recipe_builder.ModifierYAMLBuilder]] = None) → sparseml.sparsification.recipe_builder.RecipeYAMLBuilder[source] Adds a modifier group with the given name to this builder
- Parameters
name – name of new modifier group
modifier_builders – list of modifier builder objects to initialize this group with. Default is an empty list
- Returns
a reference to this object with the modifier group now added
-
build_yaml_str
() → str[source] - Returns
yaml string representation of this recipe in standard format
-
get_modifier_builders
(modifier_type: Optional[Union[Type[sparseml.optim.modifier.BaseModifier], str]] = None, modifier_groups: Optional[Union[str, List[str]]] = None)[source] - Parameters
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
modifier_groups – optional list of modifier group names to match to. Defaults to None
- Returns
all modifier builders in this recipe, filtered by type and group
-
get_modifier_group
(name: str) → Optional[List[sparseml.sparsification.recipe_builder.ModifierYAMLBuilder]][source] - Parameters
name – name of the modifier group to retrieve the modifier builders of
- Returns
reference to the list of modifier builders currently in this modifier group. if the modifier group does not exist, None will be returned
-
get_variable
(name: str, default: Optional[Any] = None) → Any[source] - Parameters
name – name of the recipe variable to return
default – default value that should be returned if the given name is not a current variable
- Returns
current value of the given variable, or the default if the variable is not set in this builder
-
has_variable
(name: str) → bool[source] - Parameters
name – name of the recipe variable to check
- Returns
True if this recipe builder has a variable with the given name. False otherwise
-
save_markdown
(file_path: str, desc: str = '')[source] Saves this recipe as a markdown file to the given path with the recipe yaml contained in the frontmatter
- Parameters
file_path – file path to save file to. if no ‘.’ character is found in the path, ‘.md’ will be added to the path
desc – optional description to add to the markdown file after the recipe YAML in the frontmatter. Default is empty string
-
save_yaml
(file_path: str)[source] Saves this recipe as a yaml file to the given path
- Parameters
file_path – file path to save file to. if no ‘.’ character is found in the path, ‘.yaml’ will be added to the path
-
set_variable
(name: str, val: Any) → sparseml.sparsification.recipe_builder.RecipeYAMLBuilder[source] Sets the given variable name to the given value
- Parameters
name – variable name to set
val – value to set the variable to
- Returns
a reference to this object with the variable now set
sparseml.sparsification.recipe_editor module
Classes for creating SparseML recipes through a series of edits based on model structure and analysis
-
class
sparseml.sparsification.recipe_editor.
MobilenetRecipeEditor
[source] Bases:
sparseml.sparsification.recipe_editor.RecipeEditor
Recipe Editor for RecipeYAMLBuilder objects with pruning modifiers for mobilenet models. Remove mobilenet depthwise layers from their pruning targets. Graphs with over 30% depthwise convs and over 30% pointwise convs are assumed to be mobilenet(-like) models
-
static
available
(model_info: sparseml.sparsification.model_info.ModelInfo, recipe_builder: sparseml.sparsification.recipe_builder.RecipeYAMLBuilder) → bool[source] this editor will be available if the builder has a pruning modifier and over 30% of the conv layers in the model are depthwise and 30% are pointwise convs
- Parameters
model_info – ModelInfo object of the model the recipe is to be created for; should contain layer information and analysis
recipe_builder – RecipeYAMLBuilder of the recipe to be edited
- Returns
True if given the inputs, this editor can edit the given recipe. False otherwise
-
static
is_depthwise_conv
(layer_info: sparseml.sparsification.model_info.LayerInfo) → bool[source]
-
static
is_pointwise_conv
(layer_info: sparseml.sparsification.model_info.LayerInfo) → bool[source]
-
static
update_recipe
(model_info: sparseml.sparsification.model_info.ModelInfo, recipe_builder: sparseml.sparsification.recipe_builder.RecipeYAMLBuilder)[source] Removes depthwise convs as pruning targets in the recipe
- Parameters
model_info – ModelInfo object of the model the recipe is to be created for; should contain layer information and analysis
recipe_builder – RecipeYAMLBuilder of the recipe to update
-
static
-
class
sparseml.sparsification.recipe_editor.
RecipeEditor
[source] Bases:
abc.ABC
Abstract class for incrementally editing a recipe
-
abstract static
available
(model_info: sparseml.sparsification.model_info.ModelInfo, recipe_builder: sparseml.sparsification.recipe_builder.RecipeYAMLBuilder) → bool[source] Abstract method to determine if this RecipeEditor is eligible to edit the given recipe
- Parameters
model_info – ModelInfo object of the model the recipe is to be created for; should contain layer information and analysis
recipe_builder – RecipeYAMLBuilder of the recipe to be edited
- Returns
True if given the inputs, this editor can edit the given recipe. False otherwise
-
abstract static
update_recipe
(model_info: sparseml.sparsification.model_info.ModelInfo, recipe_builder: sparseml.sparsification.recipe_builder.RecipeYAMLBuilder)[source] Abstract method to update a recipe given its current state as a RecipeYAMLBuilder and given the analysis in the ModelInfo object
- Parameters
model_info – ModelInfo object of the model the recipe is to be created for; should contain layer information and analysis
recipe_builder – RecipeYAMLBuilder of the recipe to update
-
abstract static
-
class
sparseml.sparsification.recipe_editor.
SkipFirstLastLayersRecipeEditor
[source] Bases:
sparseml.sparsification.recipe_editor.RecipeEditor
Recipe Editor for removing any prunable layers as pruning targets that are either the first prunable layer after model input or last prunable layer before model output. The RecipeEditor must have at least one GMPruningModifier to be available
First and last layers should have the first_prunable_layer and last_prunable_layer values set to True in their LayerInfo attributes. if no layers have this attribute then the editor will not be available
-
abstract static
available
(model_info: sparseml.sparsification.model_info.ModelInfo, recipe_builder: sparseml.sparsification.recipe_builder.RecipeYAMLBuilder) → bool[source] Available if this recipe builder has at least one GMPruningModifier and at least one layer info should have the first_prunable_layer or last_prunable_layer attribute set to True
- Parameters
model_info – ModelInfo object of the model the recipe is to be created for; should contain layer information and analysis
recipe_builder – RecipeYAMLBuilder of the recipe to be edited
- Returns
True if given the inputs, this editor can edit the given recipe. False otherwise
-
abstract static
update_recipe
(model_info: sparseml.sparsification.model_info.ModelInfo, recipe_builder: sparseml.sparsification.recipe_builder.RecipeYAMLBuilder)[source] Removes any layers from pruning that are marked as the first or last prunable layer(s)
- Parameters
model_info – ModelInfo object of the model the recipe is to be created for; should contain layer information and analysis
recipe_builder – RecipeYAMLBuilder of the recipe to update
-
abstract static
-
class
sparseml.sparsification.recipe_editor.
TieredPruningRecipeEditor
[source] Bases:
sparseml.sparsification.recipe_editor.RecipeEditor
Recipe Editor for PruningRecipeBuilder objects that have a single pruning modifier. Splits the target pruning params into three groups with target sparsities of 1.0, 0.9, and 0.8 the baseline target sparsity based on their pruning sensitivity analysis results
Baseline target sparsity can be overridden by the base_target_sparsity recipe variable. The mid and low target sparsities as a fraction of the baseline target can be overridden with the variables prune_mid_target_pct and prune_low_target_pct respectively
-
static
available
(model_info: sparseml.sparsification.model_info.ModelInfo, recipe_builder: sparseml.sparsification.recipe_builder.RecipeYAMLBuilder) → bool[source] Determines if this recipe editor can modify the recipe builder. This editor is eligible if the recipe builder is a PruningRecipeBuilder with one pruning modifier, it is in the pruning_modifiers group, and the ModelInfo includes at least one PruningSensitivityResult
- Parameters
model_info – ModelInfo object of the model the recipe is to be created for; should contain layer information and analysis
recipe_builder – RecipeYAMLBuilder of the recipe to be edited
- Returns
True if given the inputs, this editor can edit the given recipe. False otherwise
-
static
update_recipe
(model_info: sparseml.sparsification.model_info.ModelInfo, recipe_builder: sparseml.sparsification.recipe_builder.RecipeYAMLBuilder)[source] Abstract method to update a recipe given its current state as a RecipeYAMLBuilder and given the analysis in the ModelInfo object
- Parameters
model_info – ModelInfo object of the model the recipe is to be created for; should contain layer information and analysis
recipe_builder – RecipeYAMLBuilder of the recipe to update
-
static
-
sparseml.sparsification.recipe_editor.
run_avaialble_recipe_editors
(model_info: sparseml.sparsification.model_info.ModelInfo, recipe_builder: sparseml.sparsification.recipe_builder.RecipeYAMLBuilder)[source] runs all recipe editors that are available for the given model info and builder
- Parameters
model_info – ModelInfo object of the model the recipe is to be created for; should contain layer information and analysis
recipe_builder – RecipeYAMLBuilder of the recipe to update
Module contents
Functionality related to applying, describing, and supporting sparsification algorithms to models within in ML frameworks.