sparseml.optim package

Submodules

sparseml.optim.analyzer module

Code for describing layers / operators in ML framework neural networks.

class sparseml.optim.analyzer.AnalyzedLayerDesc(name: str, type_: str, params: int = 0, zeroed_params: int = 0, prunable_params: int = 0, params_dims: Optional[Dict[str, Tuple[int, …]]] = None, prunable_params_dims: Optional[Dict[str, Tuple[int, …]]] = None, execution_order: int = - 1, input_shape: Optional[Tuple[Tuple[int, …], …]] = None, output_shape: Optional[Tuple[Tuple[int, …], …]] = None, flops: int = 0, total_flops: int = 0, stride: Optional[Tuple[int, …]] = None)

Bases: object

Description of an executed neural network layer. Contains information about the number of flops, shapes, params, etc.

Parameters

• name – name of the layer

• type – type of the layer

• params – number of parameters of the layer

• zeroed_params – number of parameters with values of zero

• prunable_params – number of parameters that could be pruned

• params_dims – dimensions of parameters

• prunable_params_dims – dimensions of prunable parameters

• execution_order – execution order of the layer/operation

• input_shape – shapes of input tensors

• output_shape – shapes of output tensors

• flops – Unused

• total_flops – total number of float operations

dict() → Dict[str, Any]

Returns

A serializable dictionary representation of the current instance

static load_descs(path: str) → List

Load a list of AnalyzedLayerDesc from a json file

Parameters

path – the path to load the descriptions from

Returns

the loaded list of AnalyzedLayerDesc

static merge_descs(orig, descs: List)

Merge a layer description with a list of others

Parameters

• orig – original description

• descs – list of descriptions to merge with

Returns

a combined description

property prunable

True if the layer supports kernel sparsity (is prunable), False otherwise

Type

return

static save_descs(descs: List, path: str)

Save a list of AnalyzedLayerDesc to a json file

Parameters

• descs – a list of descriptions to save

• path – the path to save the descriptions at

property terminal

True if this is a terminal op, ie it is doing compute and is not a container, False otherwise

Type

return

sparseml.optim.learning_rate module

sparseml.optim.manager module

Code related to managers that is shared across frameworks. Managers control groups of modifiers to allow modifying the training process of a model; ex to perform model pruning.

class sparseml.optim.manager.BaseManager(modifiers: Union[List[sparseml.optim.modifier.BaseModifier], Dict[str, List[sparseml.optim.modifier.BaseModifier]]], metadata: Optional[Dict[str, Any]] = None, **kwargs)

Bases: sparseml.optim.modifier.BaseObject

Parent class meant to be used for all managers. Handles base implementations for properties and methods.

Parameters

modifiers – the modifiers to wrap

Metadata

additional (to the information provided in the recipe) data to be preserved and possibly utilized - for reproducibility and completeness

classmethod compose_staged(base_recipe: Union[str, sparseml.optim.manager.BaseManager], additional_recipe: Union[str, sparseml.optim.manager.BaseManager], keep_original_epochs: bool = False, save_path: Optional[str] = None) → sparseml.optim.manager.BaseManager

composes two recipes into a multi-stage recipe where epochs for additional_recipe are overwritten to come after base_recipe

Parameters

• base_recipe – base recipe to compose multi stage recipe with. May be a string YAML recipe, file path, or Manager object

• additional_recipe – additional recipe whose stages will be added to the base recipe. epoch ranges for additional_recipe will be adjusted to come after base_recipe unless keep_original_epochs is set. May be a string YAML recipe, file path, or Manager object

• keep_original_epochs – by default, epochs in additional_recipe will be overwritten to come after base_recipe. setting keep_original_epochs to True prevents this behavior. Default is False

• save_path – optional path string; if provided, will be used to immediately save the combined multi-stage recipe to yaml

Returns

framework Manager object with the loaded composed recipe

distillation_modifiers

epoch_modifiers

finalize_and_save_structured_modifiers(file_path: str)

saves a recipe containing only the structure modifiers of this manager. start and end epochs are overwritten so that they will be applied by epoch 0 in order

Parameters

file_path – file path to save the yaml recipe to

iter_modifiers() → Generator[None, None, sparseml.optim.modifier.BaseModifier]

Returns

generator for modifiers of this manager

learning_rate_modifiers

max_epochs

property metadata

metadata_to_string_lines(stage: Optional[str] = None) → List[str]

Parse self._metadata into list of strings. :param stage: Name of the current recipe stage.

If stage = None, we are dealing with standard, unstaged recipe.

Returns

a list of lines for a string / yaml representation of the metadata for the given stage in the manager

min_epochs

modifiers

modifiers_list_to_string_lines(modifiers: List[sparseml.optim.modifier.BaseModifier]) → List[str]

Parameters

modifiers – the modifiers to convert into string / yaml representation for within the manage

Returns

a list of lines for a string / yaml representation of the modifiers in the manager

modifiers_to_string_lines(modifiers: Union[List[sparseml.optim.modifier.BaseModifier], Dict[str, List[sparseml.optim.modifier.BaseModifier]]], include_metadata: bool = True) → List[str]

Parameters

• modifiers – the modifiers to convert into string / yaml representation for within the manage

• include_metadata – boolean indicator whether metadata shall be appended to the yaml file before saving.

Returns

a list of lines for a string / yaml representation of the modifiers in the manager

num_stages() → int

Return the number of stages of the recipe :return: number of stages

pruning_modifiers

qat_active(epoch: float) → bool

Parameters

epoch – the epoch to check if quantization aware training will be active during

Returns

True if quantization aware training will be active at the start of or within the given epoch, False otherwise

quantization_modifiers

save(file_path: str, include_metadata: bool = True)

Parameters

• file_path – the file path to save the yaml config representation to

• include_metadata – boolean indicator whether metadata shall be appended to the yaml file before saving. Default is True.

structured_modifiers

to_string_lines(include_metadata: bool = True) → List[str]

Parameters

include_metadata – boolean indicator whether metadata shall be appended to the yaml file before saving. Default is False.

Returns

a list of lines for a string / yaml representation of this instance

sparseml.optim.modifier module

Code related to modifiers that is shared across frameworks. Modifiers allow modifying the training process of a model; ex to perform model pruning.

class sparseml.optim.modifier.BaseModifier(**kwargs)

Bases: sparseml.optim.modifier.BaseObject

Parent class meant to be used for all modifiers. Handles base implementations for properties and methods.

Parameters

kwargs – standard key word args, used to support multi inheritance

static comparator(one, two) → int

Comparator implementation for Modifiers. Compares first on end_epoch, next on start_epoch, and finally on identifier.

Parameters

• one – first modifier to compare

• two – second modifier to compare

Returns

int representing where one is in relation to two

static comparator_ends(one, two) → int

Comparator implementation for Modifiers based on end_epoch. Modifiers with ends greater than another will come out higher.

Parameters

• one – first modifier to compare

• two – second modifier to compare

Returns

int representing where one is in relation to two

static comparator_identifiers(one, two) → int

Comparator implementation for Modifiers based on identifier. Modifiers with ends greater than another will come out higher.

Parameters

• one – first modifier to compare

• two – second modifier to compare

Returns

int representing where one is in relation to two

static comparator_lists(one: List[sparseml.optim.modifier.BaseModifier], two: List[sparseml.optim.modifier.BaseModifier]) → int

Comparator for list of modifiers, compares the max end, min start epochs of either lists and then the maximal identifiers of either

Parameters

• one – first list of modifiers to compare

• two – second list of modifiers to compare

Returns

int representing where one is in relation to two

static comparator_starts(one, two) → int

Comparator implementation for Modifiers based on start_epoch. Modifiers with starts greater than another will come out higher.

Parameters

• one – first modifier to compare

• two – second modifier to compare

Returns

int representing where one is in relation to two

enabled

True if the modifier is currently enabled and making updates, False otherwise

Type

return

identifier(extra: Any = '') → str

Parameters

extra – any extra identifier to append to the end of the string

Returns

generate an identifier for the current modifier based on its class name and params

initialized

static load_framework_list(yaml_str: str, framework: str)

Parameters

• yaml_str – a string representation of the yaml syntax to load modifiers

• framework – the framework to load the modifiers for

Returns

the loaded modifiers list or dictionary of stage name to stage modifiers list if given a yaml string of a staged recipe

static load_framework_obj(yaml_str: str, framework: str)

Parameters

• yaml_str – a string representation of the yaml syntax to load a modifier

• framework – the framework to load the modifier for

Returns

the loaded modifier object

props(only_serializable: bool, format_str: bool = False, format_repr: bool = False) → Dict[str, Any]

Gather all the ModifierProps for the current instance into a dictionary collection.

Parameters

• only_serializable – True if only props marked as serializable should be collected, False otherwise

• format_str – True to format the values properly for a str. Ex: None values are formatted to null and otherwise str is called

• format_repr – True to format the values properly for a repr.

Returns

the collected properties with names mapping to values

sparsification_types

static yaml_key(clazz, framework: Optional[str] = None)

create a key for loading in yaml from the class and the framework

Parameters

• clazz – the class representation to create the key for

• framework – the string representing the ML framework the modifier class is for. Default is None.

Returns

the formatted key; ex: !{framework}.{clazz.__name__}

class sparseml.optim.modifier.BaseObject(**kwargs)

Bases: object

BaseObject to accept kwargs so multiple inheritance will work with the modifier classes. kwargs param must be empty by the time this class is called.

Parameters

kwargs – standard key word args, used to support multi inheritance

class sparseml.optim.modifier.BaseProp

Bases: abc.ABC

BaseProp class meant to be implemented by any property decorators

abstract getter(func_get: Callable)

Parameters

func_get – the getter function

Returns

the recreated instance with the new getter function

abstract setter(func_set: Callable)

Parameters

func_set – the setter function

Returns

the recreated instance with the new setter function

class sparseml.optim.modifier.BaseScheduled(start_epoch: float = - 1.0, min_start: float = - 1.0, end_epoch: float = - 1.0, min_end: float = - 1.0, end_comparator: Optional[int] = 0, **kwargs)

Bases: sparseml.optim.modifier.BaseObject

Abstract class meant to be used for all scheduled modifiers. :py:func ~Modifier is also meant to be inherited alongside this class. Handles base implementations for scheduled properties and methods to allow a schedule to be enforced.

Parameters

• start_epoch – the epoch to start the scheduled modifier at

• min_start – the minimum value start_epoch can be, otherwise will raise a ValueError

• end_epoch – the epoch to end the scheduled modifier at

• min_end – the minimum value end_epoch can be, otherwise will raise a ValueError

• end_comparator – integer value representing how the end_epoch should be compared to start_epoch. if == None, then end_epoch can only be set to what its initial value was. if == -1, then end_epoch can be -1, equal to, or greater than start_epoch. if == 0, then end_epoch can be equal to or greater than start_epoch. if == 1, then end_epoch can only be greater than start_epoch.

• kwargs – standard key word args, used to support multi inheritance

end_epoch

The epoch to end the modifier at (set to -1.0 so it never ends)

Type

return

start_epoch

The epoch to start the modifier at (set to -1.0 so it starts immediately)

Type

return

validate_schedule()

Validate the schedule values of the params for the current instance are valid

class sparseml.optim.modifier.BaseUpdate(update_frequency: float, min_frequency: float, **kwargs)

Bases: sparseml.optim.modifier.BaseObject

Abstract class meant to be used for all update modifiers. :py:func ~Modifier and :py:func ~ScheduledModifier are also meant to be inherited alongside this class. Handles base implementations for scheduled properties and methods to allow updates to be enforced.

Parameters

• update_frequency – The number of epochs or fraction of epochs to update at between start and end

• min_frequency – The minimum acceptable value for update_frequency, default -1

• kwargs – standard key word args, used to support multi inheritance

update_frequency

The number of epochs or fraction of epochs to update at between start and end

Type

return

validate_update()

Validate the update schedule values of the params for the current instance are valid

class sparseml.optim.modifier.ModifierProp(serializable: bool = True, restrict_initialized: bool = True, restrict_enabled: bool = False, restrict_extras: Optional[List[str]] = None, no_serialize_val: Optional[Any] = None, func_get: Optional[Callable] = None, func_set: Optional[Callable] = None, doc: Optional[Callable] = None)

Bases: sparseml.optim.modifier.BaseProp

Property used to decorate a modifier. Use for creating getters and setters in a modifier. Handles making sure props cannot be changed after a certain point; ex after initialized. Also, marks the properties so they can be easily collected and serialized later.

Parameters

• serializable – True if the property should be serialized (ex in yaml), False otherwise. Default True

• restrict_initialized – True to keep the property from being set after initialized, False otherwise. Default True

• restrict_enabled – True to keep the property from being set after enabled, False otherwise. Default False

• restrict_extras – extra attributes to check, if any are truthy then keep from being set. Default None

• no_serialize_val – If prop is equal to this value, will not serialize the prop

• func_get – The function getter

• func_set – The function setter

• doc – The docs function

getter(func_get: Callable) → sparseml.optim.modifier.BaseProp

Create a ModifierProp based off the current instance with the getter function

Parameters

func_get – the getter function

Returns

the recreated instance with the new getter function

property no_serialize_val

a value that if the prop is equal to, will not serialize the prop

Type

return

property restrictions

The attributes to check for restricting when the attribute can be set

Type

return

property serializable

True if the property should be serialized (ex in yaml), False otherwise

Type

return

setter(func_set: Callable) → sparseml.optim.modifier.BaseProp

Create a ModifierProp based off the current instance with the setter function

Parameters

func_set – the setter function

Returns

the recreated instance with the new setter function

class sparseml.optim.modifier.ModifierYAML(framework: str)

Bases: object

A decorator to handle making a modifier class YAML ready. IE it can be loaded in through the yaml plugin easily.

Parameters

framework – the string representing the ML framework the modifier should be stored under

sparseml.optim.sensitivity module

Generic code related to sensitivity analysis.

class sparseml.optim.sensitivity.LRLossSensitivityAnalysis

Bases: object

Basic class for tracking the results from a learning rate sensitivity analysis

add_result(lr: float, loss_measurements: List[float])

dict() → Dict[str, Any]

Returns

dictionary representation of the current instance

static load_json(path: str)

Parameters

path – the path to load a previous analysis from

Returns

the analysis instance created from the json file

plot(path: Optional[str], title: Optional[str] = None) → Union[Tuple[matplotlib.figure.Figure, matplotlib.axes._axes.Axes], Tuple[None, None]]

Plot the recorded sensitivity values

Parameters

• path – the path for where to save the plot, if not supplied will display it

• title – the title of the plot to apply, defaults to ‘{plot_loss_key} LR Sensitivity’

Returns

the figure and axes if the figure was displayed; else None, None

print_res()

Print the recorded sensitivity values CSV results

property results

save_json(path: str)

Parameters

path – the path to save the json file at representing the layer sensitivities

class sparseml.optim.sensitivity.PruningLossSensitivityAnalysis

Bases: object

Analysis result for how kernel sparsity (pruning) affects the loss of a given model. Contains layer by layer results.

add_result(id_: Optional[str], name: str, index: int, sparsity: float, measurement: float, baseline: bool)

Add a result to the sensitivity analysis for a specific param

Parameters

• id – the identifier to add the result for

• name – the readable name to add the result for

• index – the index of the param as found in the model parameters

• sparsity – the sparsity to add the result for

• measurement – the loss measurement to add the result for

• baseline – True if this is a baseline measurement, False otherwise

dict() → Dict[str, Any]

Returns

dictionary representation of the current instance

static from_dict(dictionary: Dict[str, Any])

Parameters

dictionary – the dictionary to create an analysis object from

Returns

the KSLossSensitivityAnalysis instance from the json

get_result(id_or_name: str) → sparseml.optim.sensitivity.PruningSensitivityResult

get a result from the sensitivity analysis for a specific param

Parameters

id_or_name – the id or name to get the result for

Returns

the loss sensitivity results for the given id or name

static load_json(path: str)

Parameters

path – the path to load a previous analysis from

Returns

the KSLossSensitivityAnalysis instance from the json

plot(path: Optional[str], plot_integral: bool, normalize: bool = True, title: Optional[str] = None) → Union[Tuple[matplotlib.figure.Figure, matplotlib.axes._axes.Axes], Tuple[None, None]]

Parameters

• path – the path to save an img version of the chart, None to display the plot

• plot_integral – True to plot the calculated loss integrals for each layer, False to plot the averages

• normalize – normalize the values to a unit distribution (0 mean, 1 std)

• title – the title to put on the chart

Returns

the created figure and axes if path is None, otherwise (None, None)

print_res()

Print the recorded sensitivity values results

property results

the individual results for the analysis

Type

return

property results_model

the overall results for the model

Type

return

save_json(path: str)

Parameters

path – the path to save the json file at representing the layer sensitivities

class sparseml.optim.sensitivity.PruningPerfSensitivityAnalysis(num_cores: int, batch_size: int)

Bases: object

Analysis result for how kernel sparsity (pruning) affects the loss of a given model. Contains layer by layer results.

Parameters

• num_cores – number of physical cpu cores the analysis was run on

• batch_size – the input batch size the analysis was run for

add_model_result(sparsity: float, measurement: float, baseline: bool)

Add a result to the sensitivity analysis for the overall model

Parameters

• sparsity – the sparsity to add the result for

• measurement – resulting timing in seconds for the given sparsity for the measurement

• baseline – True if this is a baseline measurement, False otherwise

add_result(id_: Optional[str], name: str, index: int, sparsity: float, measurement: float, baseline: bool)

Add a result to the sensitivity analysis for a specific param

Parameters

• id – the identifier to add the result for

• name – the readable name to add the result for

• index – the index of the param as found in the model parameters

• sparsity – the sparsity to add the result for

• measurement – resulting timing in seconds for the given sparsity for the measurement

• baseline – True if this is a baseline measurement, False otherwise

property batch_size

the input batch size the analysis was run for

Type

return

dict() → Dict[str, Any]

Returns

dictionary representation of the current instance

static from_dict(dictionary: Dict[str, Any])

Parameters

dictionary – the dictionary to create an analysis object from

Returns

the KSPerfSensitivityAnalysis instance from the json

get_result(id_or_name: str) → sparseml.optim.sensitivity.PruningSensitivityResult

get a result from the sensitivity analysis for a specific param

Parameters

id_or_name – the id or name to get the result for

Returns

the loss sensitivity results for the given id or name

static load_json(path: str)

Parameters

path – the path to load a previous analysis from

Returns

the KSPerfSensitivityAnalysis instance from the json

property num_cores

number of physical cpu cores the analysis was run on

Type

return

plot(path: Optional[str], title: Optional[str] = None) → Union[Tuple[matplotlib.figure.Figure, matplotlib.axes._axes.Axes], Tuple[None, None]]

Parameters

• path – the path to save an img version of the chart, None to display the plot

• title – the title to put on the chart

Returns

the created figure and axes if path is None, otherwise (None, None)

print_res()

Print the recorded sensitivity values results

property results

the individual results for the analysis

Type

return

property results_model

the overall results for the model

Type

return

save_json(path: str)

Parameters

path – the path to save the json file at representing the layer sensitivities

class sparseml.optim.sensitivity.PruningSensitivityResult(id_: str, name: str, index: int, baseline_measurement_index: int = - 1, baseline_measurement_key: Optional[str] = None, sparse_measurements: Optional[Dict[float, List[float]]] = None)

Bases: object

A sensitivity result for a given node/param in a model. Ex: loss sensitivity or perf sensitivity

Parameters

• id – id for the node / param

• name – human readable name for the node / param

• index – index order for when the node / param is used in the model

• baseline_measurement_index – index for where the baseline measurement is stored in the sparse_measurements, if any

• sparse_measurements – the sparse measurements to prepopulate with, if any

add_measurement(sparsity: float, loss: float, baseline: bool)

add a sparse measurement to the result

Parameters

• sparsity – the sparsity the measurement was performed at

• loss – resulting loss for the given sparsity for the measurement

• baseline – True if this is a baseline measurement, False otherwise

property averages

average values of loss for each level recorded

Type

return

property baseline_average

the baseline average time to compare to for the result

Type

return

property baseline_measurement_index

index for where the baseline measurement is stored in the sparse_measurements, if any

Type

return

property baseline_measurement_key

key for where the baseline measurement is stored in the sparse_measurements, if any

Type

return

dict() → Dict[str, Any]

Returns

dictionary representation of the current instance

static from_dict(dictionary: Dict[str, Any])

Create a new loss sensitivity result from a dictionary of values. Expected to match the format as given in the dict() call.

Parameters

dictionary – the dictionary to create a result out of

Returns

the created KSLossSensitivityResult

property has_baseline

True if the result has a baseline measurement in the sparse_measurements, False otherwise

Type

return

property id_

id for the node / param

Type

return

property index

index order for when the node / param is used in the model

Type

return

property name

human readable name for the node / param

Type

return

property sparse_average

average loss across all levels recorded

Type

return

sparse_comparison(compare_index: int = - 1)

Compare the baseline average to a sparse average value through the difference: sparse - baseline

If compare_index is not given then will compare with the sparsity closest to 90%. 90% is used as a reasonable achievable baseline to keep from introducing too much noise at the extremes of the tests.

If not has_baseline, then will compare against the first index.

Parameters

compare_index – the index to compare against the baseline with, if not supplied will compare against the sparsity measurement closest to 90%

Returns

a comparison of the sparse average with the baseline (sparse - baseline)

property sparse_integral

integrated loss across all levels recorded

Type

return

property sparse_measurements

the sparse measurements

Type

return

sparseml.optim.sensitivity.default_pruning_sparsities_loss(extended: bool) → Tuple[float, …]

The default sparsities to use for checking pruning effects on the loss

Parameters

extended – extend the sparsties to return a full range instead of a subset of target sparstiies

Returns

the sparsities to check for effects on the loss

sparseml.optim.sensitivity.default_pruning_sparsities_perf() → Tuple[float, …]

Returns

the sparsities to check for effects on the loss

Module contents

Recalibration code shared across ML frameworks. Handles things like model pruning and increasing activation sparsity.