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)

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

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

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]

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)

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)

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

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]

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

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)

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)

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)

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] = [])

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]

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

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)

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

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)

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)

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

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

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)

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)

property analysis_results

list of analysis results run on this model

Type

return

abstract extract_layer_info(model: Any) → OrderedDict[str, LayerInfo]

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

Parameters

dictionary – dict serialized by dict(ModelInfo(…))

Returns

ModelInfo object created from the given dict

get_prunable_param_names() → Set[str]

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]

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

Parameters

file_path – file path to JSON file to load ModelInfo object from

Returns

the loaded ModelInfo object

save(file_path: str)

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]

Returns

dict representation of this ModelResult

static validate_model(model: Any) → Any

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)

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)

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)

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)

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]

Returns

list of sparsity values available for all model layers

get_layer_sparsity_score(layer_name: str, sparsity: float) → float

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)

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)

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)

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

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, log_types: Union[str, List[str]] = '__ALL__', **kwargs)

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

log_types: __ALL__

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

• log_types – The loggers to allow the learning rate to be logged to, default is __ALL__

corrected_lr_info(steps_per_epoch: int, start_epoch: float, end_epoch: float) → Tuple[str, Dict]

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

validate_lr_info()

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, log_types: Union[str, List[str]] = '__ALL__', end_epoch: float = - 1.0, **kwargs)

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

log_types: __ALL__

Parameters

• learning_rate – The learning rate to use once this modifier starts

• start_epoch – The epoch to start the modifier at

• log_types – The loggers to allow the learning rate to be logged to, default is __ALL__

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

validate_learning_rate()

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)

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

trainable

True if the param(s) should be made trainable, False to make them non-trainable

Type

return

validate()

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, log_types: Union[str, List[str]] = '__ALL__', **kwargs)

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

log_types: __ALL__

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

• log_types – The loggers to allow the learning rate to be logged to, default is __ALL__

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

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', log_types: Union[str, List[str]] = '__ALL__', mask_type: Union[str, List[int]] = 'unstructured', leave_enabled: bool = True, **kwargs)

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

log_types: __ALL__

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]

• log_types – The loggers to allow the learning rate to be logged to, default is __ALL__

• 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

validate()

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]

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)

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

Returns

string representation of the built Modifier as a YAML list item

copy() → sparseml.sparsification.recipe_builder.ModifierYAMLBuilder

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

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

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)

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

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

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)

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

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

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

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 = '')

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)

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

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

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

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

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

static is_pointwise_conv(layer_info: sparseml.sparsification.model_info.LayerInfo) → bool

static update_recipe(model_info: sparseml.sparsification.model_info.ModelInfo, recipe_builder: sparseml.sparsification.recipe_builder.RecipeYAMLBuilder)

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

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

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)

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

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

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)

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

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

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)

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)

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.