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

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]

type_: sparseml.sparsification.info.ModifierType

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: SparsificationInfo

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: SparsificationInfo

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 check_params_if_prunable(values)[source]

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

save(file_path: str)[source]: saves the dict representation of this ModelInfo object as a JSON file to the given file path :param file_path: file path to save to

to_dict() → Dict[str, Any][source]

Returns: dict representation of this ModelResult

static validate_model(model: Any) → Any[source]

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

class sparseml.sparsification.model_info.PruningSensitivityResultTypes(value)[source]

Bases: enum.Enum

Types of pruning sensitivity results standardized by SparseML, used as ModelResult.analysis_type

LOSS = 'pruning_sensitivity_loss'

PERF = 'pruning_sensitivity_perf'

class sparseml.sparsification.model_info.Result(*, value: Any = None, attributes: Dict[str, Any] = None)[source]

Bases: pydantic.main.BaseModel

Base class for storing the results of an analysis

attributes: Optional[Dict[str, Any]]

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: 0

end_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:

!LearningRateModifier

lr_class: ExponentialDecay

lr_kwargs:

initial_learning_rate: 0.01
decay_steps: 10000
decay_rate: 0.96

start_epoch: 0.0

end_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

validate_lr_info()[source]: Validate the values of the params for the current instance are valid

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:

!SetLearningRateModifier

start_epoch: 0.0

learning_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

validate_learning_rate()[source]

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

validate()[source]: Validate the values of the params for the current instance are valid, Should be overriden by framework specific implementations

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:

!ConstantPruningModifier

params: __ALL__
start_epoch: 0.0
end_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:

!GMPruningModifier

params: __ALL__
init_sparsity: 0.05
final_sparsity: 0.8
start_epoch: 0.0
end_epoch: 10.0
update_frequency: 1.0
inter_func: cubic
mask_type: unstructured
leave_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

validate()[source]: Validate the values of the params for the current instance are valid

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: 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)

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

build_yaml_str() → str[source]

Returns: yaml string representation of this recipe in standard format

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_builder.to_yaml_str(val: Any) → str[source]

Parameters: val – value to get yaml str value of
Returns: direct str cast of val if it is an int, float, or bool, otherwise the stripped output of yaml.dump

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

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

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

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

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.