sparseml.tensorflow_v1.optim package

Submodules

sparseml.tensorflow_v1.optim.analyzer_module module

sparseml.tensorflow_v1.optim.analyzer_module.analyze_module(session: Optional[tensorflow.python.client.session.Session], graph: Optional[tensorflow.python.framework.ops.Graph], op_names: Optional[List[str]] = None, op_types: Optional[List[str]] = None)

Analyze a module at certain layers

Parameters

• session – running session encapsulating the analyzed module

• graph – graph of the module; if None then the session is required, and the encapsulated graph is to be analyzed

• op_names – list of names of layers to be analyzed; if None then all layers are analyzed for an aggregated result

• op_types – the operation types that will be analyzed, default (Conv2D, MatMul)

Returns

the analyzed layer descriptions or the module description if no op_names

sparseml.tensorflow_v1.optim.manager module

Contains base code related to modifier managers: modifier managers handle grouping modifiers and running them together. Also handles loading modifiers from yaml files

class sparseml.tensorflow_v1.optim.manager.ScheduledModifierManager(modifiers: List[sparseml.tensorflow_v1.optim.modifier.ScheduledModifier], metadata: Optional[Dict[str, Any]] = None)

Bases: sparseml.optim.manager.BaseManager, sparseml.tensorflow_v1.optim.modifier.Modifier

The base modifier manager, handles managing multiple ScheduledModifier.

Modifiers are expected to implement up to 3 different functions for TensorFlow:

- create_ops - inject ops into the graph before the training begins

- create_extras - create extras like learning rate controls before training

- complete_graph - finalize the graph after training has completed

Life cycle:

- create model graph

- manager.create_ops()

- manager.create_extras()

- train graph

- manager.complete_graph()

- export graph

Parameters

modifiers – the modifiers to wrap

RECAL_UPDATE = 'recal_update'

complete_graph(graph: Optional[tensorflow.python.framework.ops.Graph] = None, sess: Optional[tensorflow.python.client.session.Session] = None)

Complete modifying the graph. Should be called after modifying is complete. Cleans up any ops that should be removed or reordered.

Parameters

• graph – the modified graph that should be completed and cleaned. if not supplied, then will use the default graph

• sess – the session to use for completing the modified graph. if not supplied, then will use the default session

Returns

the cleaned graph

create_ops(steps_per_epoch: int, global_step: Optional[tensorflow.python.framework.ops.Tensor] = None, graph: Optional[tensorflow.python.framework.ops.Graph] = None) → Tuple[List[Union[tensorflow.python.framework.ops.Tensor, tensorflow.python.framework.ops.Operation]], Dict[str, Any]]

Create modifying operations and tensors in the graph.

Returns a tuple containing:

- modifying ops that should be run in a session on each global step.

- named extras (ops / tensors) created in the graph that can be used

by other ops such as a learning rate for the optimizer

Parameters

• steps_per_epoch – the number of steps (batches) per training epoch

• global_step – the global step used while training. if not supplied, then will use get_or_create_global_step()

• graph – the graph to be modified, if not supplied, then will use the default graph

Returns

a tuple (list of ops, dict of named ops / tensors) to be run or used for modifying the training process

static from_yaml(file_path: Union[str, sparsezoo.objects.recipe.Recipe], add_modifiers: Optional[List[sparseml.tensorflow_v1.optim.modifier.Modifier]] = None, recipe_variables: Optional[Union[Dict[str, Any], str]] = None, metadata: Optional[Dict[str, Any]] = None)

Convenience function used to create the manager of multiple modifiers from a recipe file.

Parameters

• file_path – the path to the recipe file to load the modifier from, or a SparseZoo model stub to load a recipe for a model stored in SparseZoo. SparseZoo stubs should be preceded by ‘zoo:’, and can contain an optional ‘?recipe_type=<type>’ parameter. Can also be a SparseZoo Recipe object. i.e. ‘/path/to/local/recipe.yaml’, ‘zoo:model/stub/path’, ‘zoo:model/stub/path?recipe_type=transfer’

• add_modifiers – additional modifiers that should be added to the returned manager alongside the ones loaded from the recipe file

• recipe_variables – additional variable values to override the recipe with (i.e. num_epochs, init_lr)

Metadata

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

Returns

ScheduledModifierManager() created from the recipe file

initialize_session(sess: Optional[tensorflow.python.client.session.Session] = None)

Initialize any state for a session such as variables. This is an optional call, only needed if global_variables_initializer is not used.

Parameters

sess – the session to use for initializing

modifiers_to_string_lines(modifiers: List[sparseml.optim.modifier.BaseScheduled]) → List[str]

Parameters

modifiers – ignored and overwritten with the original (non grouped) modifiers

Returns

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

sparseml.tensorflow_v1.optim.mask_creator_pruning module

Classes for defining sparsity masks based on model parameters.

class sparseml.tensorflow_v1.optim.mask_creator_pruning.BlockPruningMaskCreator(block_shape: List[int], grouping_op_name: str = 'mean')

Bases: sparseml.tensorflow_v1.optim.mask_creator_pruning.GroupedPruningMaskCreator

Structured sparsity mask creator that groups the input tensor into blocks of shape block_shape. block_shape must divide the shape of any input tensor evenly and must have exactly 2 elements for the shape of in and out channels in the blocks.

Parameters

block_shape – The shape of blocks to strucure blocks of in and out channels in the mask by. -1 represents blocking along the entire dimension.

group_tensor(tensor: tensorflow.python.framework.ops.Tensor) → tensorflow.python.framework.ops.Tensor

Parameters

tensor – The tensor to transform

Returns

The absolute mean values of the tensor grouped by blocks of shape self._block_shape

class sparseml.tensorflow_v1.optim.mask_creator_pruning.DimensionPruningMaskCreator(dim: Union[str, int, List[int]], grouping_op_name: str = 'mean')

Bases: sparseml.tensorflow_v1.optim.mask_creator_pruning.GroupedPruningMaskCreator

Structured sparsity mask creator that groups sparsity blocks by the given dimension(s)

Parameters

dim – The index or list of indices of dimensions to group the mask by or the type of dims to prune ([‘channel’, ‘filter’])

group_tensor(tensor: tensorflow.python.framework.ops.Tensor) → tensorflow.python.framework.ops.Tensor

Parameters

tensor – The tensor to transform

Returns

The absolute mean values of the tensor grouped by the dimension(s) in self._dim

class sparseml.tensorflow_v1.optim.mask_creator_pruning.GroupedPruningMaskCreator

Bases: sparseml.tensorflow_v1.optim.mask_creator_pruning.UnstructuredPruningMaskCreator

Abstract class for a sparsity mask creator that structures masks according to grouping functions. Subclasses should implement group_tensor and _map_mask_to_tensor

create_sparsity_mask(tensor: tensorflow.python.framework.ops.Tensor, sparsity: tensorflow.python.framework.ops.Tensor) → tensorflow.python.framework.ops.Tensor

Parameters

• tensor – A tensor of a model layer’s weights

• sparsity – the target sparsity to use for assigning the masks

Returns

A sparsity mask close to the set sparsity based on the values of the input tensor

static get_grouping_op(grouping_op_name: str) → tensorflow.python.framework.ops.Operation

Parameters

grouping_op_name – name of grouping operation to get tf operation for

Returns

tf operation for grouping_op_name if available, raises error otherwise

get_mask_initializer(tensor: tensorflow.python.framework.ops.Tensor) → Callable[tensorflow.python.framework.ops.Tensor]

Parameters

tensor – A tensor of a model layer’s weights

Returns

Tensor initializer function for this sparsity mask

abstract group_tensor(tensor: tensorflow.python.framework.ops.Tensor) → tensorflow.python.framework.ops.Tensor

Parameters

tensor – The tensor to reduce in groups

Returns

The grouped tensor

class sparseml.tensorflow_v1.optim.mask_creator_pruning.PruningMaskCreator

Bases: abc.ABC

Base abstract class for a sparsity mask creator. Subclasses should define all methods for creating masks and their initializers

abstract create_sparsity_mask(tensor: tensorflow.python.framework.ops.Tensor, sparsity: tensorflow.python.framework.ops.Tensor) → tensorflow.python.framework.ops.Tensor

Parameters

• tensor – A tensor of a model layer’s weights

• sparsity – the target sparsity to use for assigning the masks

Returns

A sparsity mask close to the set sparsity based on the values of the input tensor

abstract get_mask_initializer(tensor: tensorflow.python.framework.ops.Tensor) → Callable[tensorflow.python.framework.ops.Tensor]

Parameters

tensor – A tensor of a model layer’s weights

Returns

Tensor initializer function for this sparsity mask

class sparseml.tensorflow_v1.optim.mask_creator_pruning.UnstructuredPruningMaskCreator

Bases: sparseml.tensorflow_v1.optim.mask_creator_pruning.PruningMaskCreator

Class for creating unstructured sparsity masks. Masks will be created using unstructured sparsity by pruning weights ranked by their magnitude.

create_sparsity_mask(tensor: tensorflow.python.framework.ops.Tensor, sparsity: tensorflow.python.framework.ops.Tensor) → tensorflow.python.framework.ops.Tensor

Parameters

• tensor – A tensor of a model layer’s weights

• sparsity – the target sparsity to use for assigning the masks

Returns

A sparsity mask close to the set sparsity based on the values of the input tensor

get_mask_initializer(tensor: tensorflow.python.framework.ops.Tensor) → Callable[tensorflow.python.framework.ops.Tensor]

Parameters

tensor – A tensor of a model layer’s weights

Returns

Initializer for tensor where an element is 1.0 for nonzero weights and zero for all other weights

Raise

ValueError If the dtype is not numeric or boolean

sparseml.tensorflow_v1.optim.mask_creator_pruning.load_mask_creator(obj: Union[str, Iterable[int]]) → sparseml.tensorflow_v1.optim.mask_creator_pruning.PruningMaskCreator

Parameters

obj – Formatted string or iterable of block_shape specifying SparsityMaskCreator object to return

Returns

SparsityMaskCreator object created from obj

sparseml.tensorflow_v1.optim.mask_pruning module

Code related to applying a mask onto a variable to impose kernel sparsity, aka model pruning, on a TensorFlow graph.

class sparseml.tensorflow_v1.optim.mask_pruning.PruningOpVars(op, op_input, update, mask, masked)

Bases: tuple

property mask

Alias for field number 3

property masked

Alias for field number 4

property op

Alias for field number 0

property op_input

Alias for field number 1

property update

Alias for field number 2

class sparseml.tensorflow_v1.optim.mask_pruning.PruningScope

Bases: object

Convenience class for dealing with scope and names for kernel sparsity in the tf graph.

NM_KS = 'nm_ks'

NM_KS_OPS = 'nm_ks_ops'

OPS = 'ops'

OPS_INPUT = 'input_ops'

OPS_SCHEDULE = 'schedule_ops'

OPS_SPARSITY = 'sparsity_ops'

OPS_SUMMARY = 'summary_ops'

OPS_UPDATE = 'update_ops'

OP_COND_UPDATE = 'nm_conditional_update'

OP_MASKED_VAR = 'nm_masked_var'

OP_MASK_ASSIGN = 'nm_mask_assign'

OP_MASK_UPDATE = 'nm_mask_update'

OP_MASK_UPDATE_NO_OP = 'nm_mask_update_no_op'

OP_PRUNE_VARS_ASSIGN = 'nm_prune_vars_assign'

OP_SAVE = 'nm_save'

OP_SPARSITY = 'nm_sparsity'

OP_UPDATE_READY = 'nm_update_ready'

OP_WEIGHT_UPDATE = 'nm_weight_update'

VAR_MASK = 'nm_mask'

VAR_THRESHOLD = 'nm_threshold'

static collection_name(ks_group: str, name: str) → str

Create a predictable name for a given variable / op in a group for lookup / storage in a collection

Parameters

• ks_group – the group identifier the name belongs under

• name – the name of the op or variable to be stored or retrieved

Returns

the formatted name for use in a collection

static general(ks_group: str, additional: Optional[str] = None, trailing_slash: bool = False)

Create a general kernel sparsity scope in the tf graph. Use cases are for generic ops like target sparsity, conditional updates, etc.

Parameters

• ks_group – the group identifier the scope should be created under

• additional – any additional scope that should be added to the end

• trailing_slash – include a trailing forward slash if True, else False

Returns

the proper scope

static model(op_tens: tensorflow.python.framework.ops.Tensor, ks_group: str, additional: Optional[str] = None, trailing_slash: bool = False) → str

Create a model specific kernel sparsity scope in the tf graph. Use cases are for the specific mask, threshold, etc variables to induce sparsity along with the ops to update those vars.

Parameters

• op_tens – the op tensor to create the scope for

• ks_group – the group identifier the scope should be created under

• additional – any additional scope that should be added to the end

• trailing_slash – include a trailing forward slash if True, else False

Returns

the proper scope

sparseml.tensorflow_v1.optim.mask_pruning.apply_op_vars_masks(pruning_op_vars: List[sparseml.tensorflow_v1.optim.mask_pruning.PruningOpVars], ks_group: str, sess: tensorflow.python.client.session.Session)

Apply the masks to the original ops input var so that it can be saved with the desired sparsity for later.

Parameters

• pruning_op_vars – the list of named tuples containing the sparse mask and the op variable to apply the sparse mask to

• ks_group – the group to create the assign ops under

• sess – the session to use to run the assign

sparseml.tensorflow_v1.optim.mask_pruning.create_graph_ops_pruning(graph: tensorflow.python.framework.ops.Graph, var_names: List[str], sparsity: tensorflow.python.framework.ops.Tensor, update_ready: tensorflow.python.framework.ops.Tensor, leave_enabled: bool, is_after_end_step: tensorflow.python.framework.ops.Tensor, ks_group: str, mask_creator: sparseml.tensorflow_v1.optim.mask_creator_pruning.PruningMaskCreator) → List[sparseml.tensorflow_v1.optim.mask_pruning.PruningOpVars]

Creates the necessary variables and operators to gradually apply sparsity to a given list of operators in a graph.

Handles setting a mask on an operator to the given sparsity. Sets the mask based on pruning away the lowest absolute magnitude weights.

Parameters

• graph – the tf graph to pull the operator out of for applying the pruning to

• var_names – the names or regex patterns of names of variables to prune in the graph to the given sparsity

• sparsity – the target sparsity to use for assigning the masks

• update_ready – the tensor where if true will update the mask from sparsity, if false will not update the mask

• leave_enabled – True to continue masking the weights after end_epoch, False to stop masking

• is_after_end_step – tensor that is true if the current global step is after end_epoch

• ks_group – the group identifier the scope should be created under

• mask_creator – optional object to define sparisty mask creation

Returns

a list of the created named tuples each containing the assignment op, mask variable, threshold tensor, and masked tensor

sparseml.tensorflow_v1.optim.mask_pruning.create_ks_schedule_ops(global_step: tensorflow.python.ops.variables.VariableV1, begin_step: int, end_step: int, update_step_freq: int, init_sparsity: float, final_sparsity: float, exponent: float, ks_group: str) → Tuple[tensorflow.python.framework.ops.Tensor, tensorflow.python.framework.ops.Tensor]

Create a gradual schedule for model pruning (kernel sparsity). Creates a sparsity tensor that goes from init_sparsity til final_sparsity starting at begin_step and ending at end_step. Uses the global_step to map those. Additionally creates an update_ready tensor that is True if an update to the sparsity tensor should be run, False otherwise.

Parameters

• global_step – the global optimizer step for the training graph

• begin_step – the global step to begin pruning at

• end_step – the global step to end pruning at

• update_step_freq – the number of global steps between each weight update

• init_sparsity – the starting value for sparsity of a weight tensor to be enforce

• final_sparsity – the end value for sparsity for a weight tensor to be enforce

• exponent – the exponent to use for interpolating between init_sparsity and final_sparsity higher values will lead to larger sparsity steps at the beginning vs the end ie: linear (1) vs cubic (3)

• ks_group – the group identifier the scope should be created under

Returns

a tuple containing the signal for update_ready and the target sparsity

sparseml.tensorflow_v1.optim.mask_pruning.create_ks_scheduled_constant_graph_ops(graph: tensorflow.python.framework.ops.Graph, global_step: tensorflow.python.ops.variables.VariableV1, var_names: List[str], begin_step: int, end_step: int, ks_group: str) → Tuple[tensorflow.python.framework.ops.Tensor, List[sparseml.tensorflow_v1.optim.mask_pruning.PruningOpVars]]

Creates constant model pruning ops. Does not modify the graph.

Parameters

• graph – the tf graph to pull the operator out of for applying the pruning to

• global_step – the global optimizer step for the training graph

• var_names – a list of names or regex patterns to create constant ops for within the graph

• begin_step – the global step to begin pruning at

• end_step – the global step to end pruning at

• ks_group – the group identifier the scope should be created under

Returns

a tuple containing the update operation to run in a session, a list of the pruning ops and vars for each desired op in the graph

sparseml.tensorflow_v1.optim.mask_pruning.create_op_pruning(op: tensorflow.python.framework.ops.Operation, op_input: tensorflow.python.framework.ops.Tensor, sparsity: tensorflow.python.framework.ops.Tensor, update_ready: tensorflow.python.framework.ops.Tensor, leave_enabled: bool, is_after_end_step: tensorflow.python.framework.ops.Tensor, ks_group: str, mask_creator: sparseml.tensorflow_v1.optim.mask_creator_pruning.PruningMaskCreator) → sparseml.tensorflow_v1.optim.mask_pruning.PruningOpVars

Creates the necessary variables and operators to gradually apply sparsity to an operators variable.

Handles setting a mask on an operator to the given sparsity. Sets the mask based on pruning away the lowest absolute magnitude weights.

Parameters

• op – the operation to prune to the given sparsity

• op_input – the variable of the parameter within op to prune

• sparsity – the target sparsity to use for assigning the masks

• update_ready – the tensor where if true will update the mask from sparsity, if false will not update the mask

• leave_enabled – True to continue masking the weights after end_epoch, False to stop masking

• is_after_end_step – tensor that is true if the current global step is after end_epoch

• ks_group – the group identifier the scope should be created under

• mask_creator – object to define sparisty mask creation

Returns

a named tuple containing the assignment op, mask variable, threshold tensor, and masked tensor

sparseml.tensorflow_v1.optim.mask_pruning.create_summaries_pruning(pruning_op_vars: List[sparseml.tensorflow_v1.optim.mask_pruning.PruningOpVars])

Create TensorBoard summary ops in the current graph for the given list of PruningOpVars.

Parameters

pruning_op_vars – the list of named tuples containing the masked input to the pruned op to record sparsity for in TensorBoard.

Returns

the created summaries for the pruned op vars

sparseml.tensorflow_v1.optim.mask_pruning.get_or_create_graph_ops_pruning(graph: tensorflow.python.framework.ops.Graph, var_names: List[str], sparsity: tensorflow.python.framework.ops.Tensor, update_ready: tensorflow.python.framework.ops.Tensor, leave_enabled: bool, is_after_end_step: tensorflow.python.framework.ops.Tensor, ks_group: str, mask_creator: sparseml.tensorflow_v1.optim.mask_creator_pruning.PruningMaskCreator) → List[sparseml.tensorflow_v1.optim.mask_pruning.PruningOpVars]

Creates or retrieves (if previously created) the necessary variables and operators to gradually apply sparsity to a given list of operators in a graph.

Handles setting a mask on an operator to the given sparsity. Sets the mask based on pruning away the lowest absolute magnitude weights.

Parameters

• graph – the tf graph to pull the operator out of for applying the pruning to

• var_names – the names or regex patterns of names of variables to prune in the graph to the given sparsity

• sparsity – the target sparsity to use for assigning the masks

• update_ready – the tensor where if true will update the mask from sparsity, if false will not update the mask

• leave_enabled – True to continue masking the weights after end_epoch, False to stop masking

• is_after_end_step – tensor that is true if the current global step is after end_epoch

• ks_group – the group identifier the scope should be created under

• mask_creator – optional object to define sparisty mask creation

Returns

a list of the created or retrieved named tuples each containing the assignment op, mask variable, threshold tensor, and masked tensor

sparseml.tensorflow_v1.optim.mask_pruning.get_or_create_ks_schedule_ops(global_step: tensorflow.python.framework.ops.Tensor, begin_step: int, end_step: int, update_step_freq: int, init_sparsity: float, final_sparsity: float, exponent: float, ks_group: str) → Tuple[tensorflow.python.framework.ops.Tensor, tensorflow.python.framework.ops.Tensor]

Creates or retrieves (if previously created) a gradual schedule for model pruning (kernel sparsity). Creates a sparsity tensor that goes from init_sparsity til final_sparsity starting at begin_step and ending at end_step. Uses the global_step to map those. Additionally creates an update_ready tensor that is True if an update to the sparsity tensor should be run, False otherwise.

Parameters

• global_step – the global optimizer step for the training graph

• begin_step – the global step to begin pruning at

• end_step – the global step to end pruning at

• update_step_freq – the number of global steps between each weight update

• init_sparsity – the starting value for sparsity of a weight tensor to be enforce

• final_sparsity – the end value for sparsity for a weight tensor to be enforce

• exponent – the exponent to use for interpolating between init_sparsity and final_sparsity higher values will lead to larger sparsity steps at the beginning vs the end ie: linear (1) vs cubic (3)

• ks_group – the group identifier the scope should be created under

Returns

a tuple containing the signal for update_ready and the target sparsity

sparseml.tensorflow_v1.optim.mask_pruning.get_or_create_ks_scheduled_graph_ops(graph: tensorflow.python.framework.ops.Graph, global_step: tensorflow.python.ops.variables.VariableV1, var_names: List[str], begin_step: int, end_step: int, update_step_freq: int, init_sparsity: float, final_sparsity: float, exponent: float, leave_enabled: bool, ks_group: str, mask_creator: sparseml.tensorflow_v1.optim.mask_creator_pruning.PruningMaskCreator) → Tuple[tensorflow.python.framework.ops.Tensor, List[sparseml.tensorflow_v1.optim.mask_pruning.PruningOpVars], tensorflow.python.framework.ops.Tensor, tensorflow.python.framework.ops.Tensor]

Gets or creates model pruning (kernel sparsity) ops and vars in the graph to be applied over a specific schedule. Creates them for the var_names in the graph such that they follow a schedule from begin_step to end_step starting at init_sparsity and ending at final_sparsity.

Parameters

• graph – the tf graph to pull the operator out of for applying the pruning to

• global_step – the global optimizer step for the training graph

• var_names – the names or regex patterns of names of variables to prune in the graph

• begin_step – the global step to begin pruning at

• end_step – the global step to end pruning at

• update_step_freq – the number of global steps between each weight update

• init_sparsity – the starting value for sparsity of a weight tensor to be enforce

• final_sparsity – the end value for sparsity for a weight tensor to be enforce

• exponent – the exponent to use for interpolating between init_sparsity and final_sparsity higher values will lead to larger sparsity steps at the beginning vs the end ie: linear (1) vs cubic (3)

• leave_enabled – True to continue masking the weights after end_epoch, False to stop masking

• ks_group – the group identifier the scope should be created under

• mask_creator – optional object to define sparisty mask creation

Returns

a tuple containing the update operation to run in a session, a list of the pruning ops and vars for each desired op in the graph, the tensor containing the update_ready signal for the pruning ops, the tensor containing the set sparsity for the pruning ops

sparseml.tensorflow_v1.optim.modifier module

Contains base code related to modifiers: objects that modify some aspect of the training process for a model. For example, learning rate schedules or kernel sparsity (weight pruning) are implemented as modifiers.

class sparseml.tensorflow_v1.optim.modifier.Modifier(**kwargs)

Bases: sparseml.optim.modifier.BaseModifier

Base modifier class that all TensorFlow modifiers should derive themselves from. Handles setting up the expected contracts for modifying graphs, ops, and extras.

Modifiers are expected to implement up to 3 different functions for TensorFlow:

- create_ops - inject ops into the graph before the training begins

- create_extras - create extras like learning rate controls before training

- complete_graph - finalize the graph after training has completed

Life cycle:

- create model graph

- manager.create_ops()

- manager.create_extras()

- train graph

- manager.complete_graph()

- export graph

Parameters

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

complete_graph(graph: tensorflow.python.framework.ops.Graph, sess: tensorflow.python.client.session.Session)

Complete modifying the graph. Should be called after modifying is complete. Cleans up any ops that should be removed or reordered.

Parameters

• graph – the modified graph that should be completed and cleaned

• sess – the session to use for completing the modified graph

Returns

the cleaned graph

create_ops(steps_per_epoch: int, global_step: tensorflow.python.framework.ops.Tensor, graph: tensorflow.python.framework.ops.Graph) → Tuple[List[Union[tensorflow.python.framework.ops.Tensor, tensorflow.python.framework.ops.Operation]], Dict[str, Any]]

Create modifying operations and tensors in the graph.

Returns a tuple containing:

- modifying ops that should be run in a session on each global step.

- named extras (ops / tensors) created in the graph that can be used

by other ops such as a learning rate for the optimizer

Parameters

• steps_per_epoch – the number of steps (batches) per training epoch

• global_step – the global step used while training

• graph – the graph to be modified

Returns

a tuple (list of ops, dict of named ops / tensors) to be run or used for modifying the training process

get_group() → Any

Function to be override by a subclass indicating the modifier container into which the subclass should be combined As an example, the two learning rate modifier classes SetLearningRateModifier and LearningRateModifier return GroupLearningRateModifier, meaning that a sequence of those LR modifier instances are grouped into the GroupLearningRateModifier, which is where the final learning rate is computed

initialize_session(sess: tensorflow.python.client.session.Session)

Initialize any state for a session such as variables.

Parameters

sess – the session to use for initializing

static load_list(yaml_str: str)

Parameters

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

Returns

the loaded modifiers list

static load_obj(yaml_str: str)

Parameters

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

Returns

the loaded modifier object

modify_estimator(estimator: tensorflow_estimator.python.estimator.estimator.Estimator, steps_per_epoch: int)

Modify a tf Estimator. Overrides the model_fn so that on invocation it creates the original graph and then calls into create_ops. Additionally will recreate the Scaffold with a new Save instance to save all variables in the modified graph.

Note, learning_rate and other specific tensors that needed to be retrieved from the extras in create_ops and passed to another implementation will not work with this flow.

Parameters

• estimator – the tf Estimator to modify

• steps_per_epoch – number of steps per training epoch

class sparseml.tensorflow_v1.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.tensorflow_v1.optim.modifier.ModifierSessionRunHook(mod_ops: List[Union[tensorflow.python.framework.ops.Tensor, tensorflow.python.framework.ops.Operation]])

Bases: tensorflow.python.training.session_run_hook.SessionRunHook

A session run hook for the tf Estimator flow. Used to integrate so that any extra ops for modifying the graph can be executed each on each step of the estimator training process.

Parameters

• modifier – the modifier to run the hook for

• steps_per_epoch – number of steps (or batches) taken per epoch

• mod_ops – the ops returned from calling create_ops on the modifier

• mod_extras – the extras returned from calling create_ops on the modifier

after_run(run_context, run_values)

Called by the estimator after each call to run()

Parameters

• run_context – run_context passed in during training

• run_values – a SessionRunValues object passed in during training

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

Bases: sparseml.tensorflow_v1.optim.modifier.Modifier, sparseml.optim.modifier.BaseScheduled

The base scheduled update modifier implementation, all scheduled modifiers should inherit from this class. Offers convenient properties needed for scheduled update modifiers: start_epoch, end_epoch

Modifiers are expected to implement up to 3 different functions for TensorFlow:

- create_ops - inject ops into the graph before the training begins

- create_extras - create extras like learning rate controls before training

- complete_graph - finalize the graph after training has completed

Life cycle:

- create model graph

- manager.create_ops()

- manager.create_extras()

- train graph

- manager.complete_graph()

- export graph

Parameters

• start_epoch – The epoch to start the modifier at

• end_epoch – The epoch to end the modifier at

• min_start – The minimum acceptable value for start_epoch, default -1

• min_end – The minimum acceptable value for end_epoch, default 0

• 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, 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

start_end_steps(steps_per_epoch: int, after_optim: bool) → Tuple[int, int]

Calculate the start and end steps for this modifier given a certain amount of steps per epoch

Parameters

• steps_per_epoch – the number of steps (or batches) taken per epoch

• after_optim – True if the start and end are for an operation after the optimizer update step has run, False for before

Returns

a tuple containing (the converted start step, the converted end step)

class sparseml.tensorflow_v1.optim.modifier.ScheduledUpdateModifier(start_epoch: float = - 1.0, end_epoch: float = - 1.0, min_start: float = - 1.0, min_end: float = - 1.0, end_comparator: int = 0, update_frequency: float = - 1.0, min_frequency: float = - 1.0, **kwargs)

Bases: sparseml.tensorflow_v1.optim.modifier.ScheduledModifier, sparseml.optim.modifier.BaseUpdate

The base scheduled update modifier implementation, all scheduled update modifiers should inherit from this class. Offers convenient properties needed for scheduled update modifiers: update_frequency

Modifiers are expected to implement up to 3 different functions for TensorFlow:

- create_ops - inject ops into the graph before the training begins

- create_extras - create extras like learning rate controls before training

- complete_graph - finalize the graph after training has completed

Life cycle:

- create model graph

- manager.create_ops()

- manager.create_extras()

- train graph

- manager.complete_graph()

- export graph

Parameters

• start_epoch – The epoch to start the modifier at

• end_epoch – The epoch to end the modifier at

• min_start – The minimum acceptable value for start_epoch, default -1

• min_end – The minimum acceptable value for end_epoch, default 0

• end_comparator – integer value representing how the end_epoch should be compared to start_epoch. if == -1, then end_epoch can be -1, equal, 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.

• 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_steps(steps_per_epoch: int) → int

Calculate the update frequency steps for this modifier given a certain amount of steps per epoch

Parameters

steps_per_epoch – the number of steps (or batches) taken per epoch

Returns

a tuple containing (the converted start step, the converted end step)

class sparseml.tensorflow_v1.optim.modifier.TensorFlowModifierYAML

Bases: sparseml.optim.modifier.ModifierYAML

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

sparseml.tensorflow_v1.optim.modifier_epoch module

Contains code for epoch modifiers in tensorflow_v1

class sparseml.tensorflow_v1.optim.modifier_epoch.EpochRangeModifier(start_epoch: float, end_epoch: float)

Bases: sparseml.sparsification.modifier_epoch.EpochRangeModifier, sparseml.tensorflow_v1.optim.modifier.ScheduledModifier

Simple modifier to set the range of epochs to train over for the recalibration process. 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

sparseml.tensorflow_v1.optim.modifier_lr module

Learning rate modifiers for TensorFlow models

class sparseml.tensorflow_v1.optim.modifier_lr.GroupLearningRateModifier(lr_modifiers: List[sparseml.tensorflow_v1.optim.modifier.ScheduledModifier])

Bases: sparseml.tensorflow_v1.optim.modifier.ScheduledModifier

Combining multiple LR modifiers, correctly compute the learning rate

Parameters

lr_modifiers – List of LR modifiers to combine

create_ops(steps_per_epoch: int, global_step: tensorflow.python.ops.variables.VariableV1, graph: tensorflow.python.framework.ops.Graph) → Tuple[List[Union[tensorflow.python.framework.ops.Tensor, tensorflow.python.framework.ops.Operation]], Dict[str, Any]]

Create switch case computing the learning rate at a given global step and extras created by individual LR modifiers

Parameters

• steps_per_epoch – the number of steps per training epoch

• global_step – the global step used while training

• graph – the graph to be modified

Returns

a tuple (list of empty ops, dict of named ops/tensors for learning rate and summaries as extras)

class sparseml.tensorflow_v1.optim.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)

Bases: sparseml.sparsification.modifier_lr.LearningRateModifier, sparseml.tensorflow_v1.optim.modifier.ScheduledUpdateModifier

Modifier to set the learning rate to follow specific schedulers within a period of epochs. The following schedulers are current supported: ExponentialLR, StepLR, MultiStepLR

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

create_ops(steps_per_epoch: int, global_step: Optional[tensorflow.python.ops.variables.VariableV1], graph: Optional[tensorflow.python.framework.ops.Graph]) → Tuple[List[Union[tensorflow.python.framework.ops.Tensor, tensorflow.python.framework.ops.Operation]], Dict[str, Any]]

Create ops to update the learning rate at the current global step

Parameters

• steps_per_epoch – the number of steps (batches) per training epoch

• global_step – the global step used while training

• graph – the graph to be modified

Returns

a tuple (empty list of ops, dict of learning rate and summaries)

get_group()

Returns

The group modifier class into which this modifier needs to be combined

class sparseml.tensorflow_v1.optim.modifier_lr.SetLearningRateModifier(learning_rate: float, start_epoch: float = - 1, end_epoch: float = - 1)

Bases: sparseml.sparsification.modifier_lr.SetLearningRateModifier, sparseml.tensorflow_v1.optim.modifier.ScheduledModifier

Modifier to set the learning rate to a specific value at a certain point in the training process. Once that point is reached, will update the optimizer’s params with the learning rate

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

• end_epoch – unused and should not be set

create_ops(steps_per_epoch: int, global_step: Optional[tensorflow.python.ops.variables.VariableV1], graph: Optional[tensorflow.python.framework.ops.Graph]) → Tuple[List[Union[tensorflow.python.framework.ops.Tensor, tensorflow.python.framework.ops.Operation]], Dict[str, Any]]

Create ops to set the learning rate at a given value if the global step reaches a given value

Parameters

• steps_per_epoch – the number of steps (batches) per training epoch

• global_step – the global step used while training

• graph – the graph to be modified

Returns

a tuple (empty list of ops, dict of learning rate and logging summaries)

get_group()

Returns

The group modifier class into which this modifier needs to be combined

sparseml.tensorflow_v1.optim.modifier_params module

Modifier for changing the state of a modules params while training according to certain update formulas or patterns.

class sparseml.tensorflow_v1.optim.modifier_params.TrainableParamsModifier(params: Union[str, List[str]], trainable: bool, params_strict: bool = True, start_epoch: float = - 1.0, end_epoch: float = - 1.0)

Bases: sparseml.sparsification.modifier_params.TrainableParamsModifier, sparseml.tensorflow_v1.optim.modifier.ScheduledModifier

Modifier to control the params for a given list of parameters. Applies the trainability over all epochs. 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

complete_graph(graph: tensorflow.python.framework.ops.Graph, sess: tensorflow.python.client.session.Session)

Complete modifying the graph. Resets the objects filtered variables to their original trainability

Parameters

• graph – the modified graph that should be completed and cleaned. if not supplied, then will use the default graph

• sess – the session to use for completing the modified graph. if not supplied, then will use the default session

create_ops(steps_per_epoch: int, global_step: tensorflow.python.framework.ops.Tensor, graph: tensorflow.python.framework.ops.Graph) → Tuple[List[Union[tensorflow.python.framework.ops.Tensor, tensorflow.python.framework.ops.Operation]], Dict[str, Any]]

Create the sparsity ops to modify the training graph according to the settings for the current instance.

Parameters

• steps_per_epoch – the number of steps (batches) per training epoch

• global_step – the global step used while training

• graph – the graph to be modified

Returns

a tuple (list of ops, dict of named ops / tensors) to be run or used for modifying the training process.

validate()

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

sparseml.tensorflow_v1.optim.modifier_pruning module

Modifiers for inducing / enforcing kernel sparsity (model pruning) on models while pruning.

class sparseml.tensorflow_v1.optim.modifier_pruning.ConstantPruningModifier(params: Union[str, List[str]], start_epoch: float = - 1, end_epoch: float = - 1)

Bases: sparseml.sparsification.modifier_pruning.ConstantPruningModifier, sparseml.tensorflow_v1.optim.modifier.ScheduledModifier

Holds 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 pruning 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

complete_graph(graph: tensorflow.python.framework.ops.Graph, sess: tensorflow.python.client.session.Session)

Complete modifying the graph. Resets the pruned op’s variables using the created masks to zero out the pruned weights for saving.

Parameters

• graph – the modified graph that should be completed and cleaned. if not supplied, then will use the default graph

• sess – the session to use for completing the modified graph. if not supplied, then will use the default session

Returns

the cleaned graph

create_ops(steps_per_epoch: int, global_step: tensorflow.python.framework.ops.Tensor, graph: tensorflow.python.framework.ops.Graph) → Tuple[List[Union[tensorflow.python.framework.ops.Tensor, tensorflow.python.framework.ops.Operation]], Dict[str, Any]]

Create the sparsity ops to modify the training graph according to the settings for the current instance.

Parameters

• steps_per_epoch – the number of steps (batches) per training epoch

• global_step – the global step used while training

• graph – the graph to be modified

Returns

a tuple (list of ops, dict of named ops / tensors) to be run or used for modifying the training process.

initialize_session(sess: tensorflow.python.client.session.Session)

Initialize the mask variables for pruning.

Parameters

sess – the session to use for initializing

ks_group

property prune_op_vars

the created pruning op vars in the graph if create_ops has been called, else None

Type

return

property sparsity

the created sparsity tensor for setting the pruning ops if create_ops has been called, else None

Type

return

property update_ready

the created update_ready tensor for setting the pruning ops if create_ops has been called, else None

Type

return

class sparseml.tensorflow_v1.optim.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], sparseml.tensorflow_v1.optim.mask_creator_pruning.PruningMaskCreator] = 'unstructured', leave_enabled: bool = True)

Bases: sparseml.sparsification.modifier_pruning.GMPruningModifier, sparseml.tensorflow_v1.optim.modifier.ScheduledUpdateModifier

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 pruning 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, or a SparsityMaskCreator object. 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

complete_graph(graph: tensorflow.python.framework.ops.Graph, sess: tensorflow.python.client.session.Session)

Complete modifying the graph. Resets the pruned op’s variables using the created masks to zero out the pruned weights for saving.

Parameters

• graph – the modified graph that should be completed and cleaned. if not supplied, then will use the default graph

• sess – the session to use for completing the modified graph. if not supplied, then will use the default session

Returns

the cleaned graph

create_ops(steps_per_epoch: int, global_step: tensorflow.python.framework.ops.Tensor, graph: tensorflow.python.framework.ops.Graph) → Tuple[List[Union[tensorflow.python.framework.ops.Tensor, tensorflow.python.framework.ops.Operation]], Dict[str, Any]]

Create the sparsity ops to modify the training graph according to the settings for the current instance.

Parameters

• steps_per_epoch – the number of steps (batches) per training epoch

• global_step – the global step used while training

• graph – the graph to be modified

Returns

a tuple (list of ops, dict of named ops / tensors) to be run or used for modifying the training process.

exponent

initialize_session(sess: tensorflow.python.client.session.Session)

Initialize the mask variables for pruning.

Parameters

sess – the session to use for initializing

ks_group

property prune_op_vars

the created pruning op vars in the graph if create_ops has been called, else None

Type

return

property sparsity

the created sparsity tensor for setting the pruning ops if create_ops has been called, else None

Type

return

property update_ready

the created update_ready tensor for setting the pruning ops if create_ops has been called, else None

Type

return

sparseml.tensorflow_v1.optim.schedule_lr module

Learning rate schedules implementations for TensorFlow

sparseml.tensorflow_v1.optim.schedule_lr.multi_step_lr_schedule(global_step: tensorflow.python.framework.ops.Tensor, start_step: int, milestone_steps: List[int], init_lr: float, gamma: float, name: str = 'multi_step_lr_schedule')

Create a multi step learning rate schedule in the current graph. Multiplies init_lr by gamma after each milestone has passed. Ex: lr = init_lr * (gamma ** NUM_UPDATES)

Parameters

• global_step – the global step used for training

• start_step – the step to start the exponential schedule on

• milestone_steps – a list of steps to decrease the learning rate at, these are the number of steps that must pass after start_step to decrease lr

• init_lr – the learning rate to start the schedule with

• gamma – the decay weight to decrease init_lr by after every step_size interval

• name – the name scope to create the graph under

Returns

the calculated learning rate tensor

sparseml.tensorflow_v1.optim.schedule_lr.step_lr_schedule(global_step: tensorflow.python.framework.ops.Tensor, start_step: int, end_step: int, step_size: int, init_lr: float, gamma: float, name: str = 'exponential_lr_schedule') → tensorflow.python.framework.ops.Tensor

Create an exponential learning rate schedule in the current graph. Multiplies init_lr by gamma after each step_size interval has passed. Ex: lr = init_lr * (gamma ** NUM_UPDATES)

Parameters

• global_step – the global step used for training

• start_step – the step to start the exponential schedule on

• end_step – the step to end the exponential schedule on, can be set to -1 and in that event will continually update the LR

• step_size – the number of steps between each gamma update to the init_lr

• init_lr – the learning rate to start the schedule with

• gamma – the decay weight to decrease init_lr by after every step_size interval

• name – the name scope to create the graph under

Returns

the calculated learning rate tensor

sparseml.tensorflow_v1.optim.sensitivity_pruning module

Sensitivity analysis implementations for kernel sparsity on Graphs against loss funcs.

class sparseml.tensorflow_v1.optim.sensitivity_pruning.SparsePruningOpVars(op_vars, sparsity)

Bases: tuple

property op_vars

Alias for field number 0

property sparsity

Alias for field number 1

sparseml.tensorflow_v1.optim.sensitivity_pruning.pruning_loss_sens_magnitude(graph: Optional[tensorflow.python.framework.ops.Graph] = None, sess: Optional[tensorflow.python.client.session.Session] = None, sparsity_levels: Union[List[float], Tuple[float, …]] = (0.0, 0.01, 0.02, 0.03, 0.04, 0.05, 0.06, 0.07, 0.08, 0.09, 0.1, 0.11, 0.12, 0.13, 0.14, 0.15, 0.16, 0.17, 0.18, 0.19, 0.2, 0.21, 0.22, 0.23, 0.24, 0.25, 0.26, 0.27, 0.28, 0.29, 0.3, 0.31, 0.32, 0.33, 0.34, 0.35, 0.36, 0.37, 0.38, 0.39, 0.4, 0.41, 0.42, 0.43, 0.44, 0.45, 0.46, 0.47, 0.48, 0.49, 0.5, 0.51, 0.52, 0.53, 0.54, 0.55, 0.56, 0.57, 0.58, 0.59, 0.6, 0.61, 0.62, 0.63, 0.64, 0.65, 0.66, 0.67, 0.68, 0.69, 0.7, 0.71, 0.72, 0.73, 0.74, 0.75, 0.76, 0.77, 0.78, 0.79, 0.8, 0.81, 0.82, 0.83, 0.84, 0.85, 0.86, 0.87, 0.88, 0.89, 0.9, 0.91, 0.92, 0.93, 0.94, 0.95, 0.96, 0.97, 0.98, 0.99)) → sparseml.optim.sensitivity.PruningLossSensitivityAnalysis

Approximated kernel sparsity (pruning) loss analysis for a given model. Returns the results for each prunable param (conv, linear) in the model. Approximated by taking the magnitudes of the weights.

Parameters

• graph – the graph to inject pruning ops and vars into, if not supplied uses get_default_graph()

• sess – the session to use

• sparsity_levels – the sparsity levels to calculate the loss for for each param

Returns

the analysis results for the model

sparseml.tensorflow_v1.optim.sensitivity_pruning.pruning_loss_sens_one_shot(op_vars: List[sparseml.tensorflow_v1.optim.sensitivity_pruning.SparsePruningOpVars], loss_tensor: tensorflow.python.framework.ops.Tensor, steps_per_measurement: int, add_ops_creator: Optional[Callable[int, List[tensorflow.python.framework.ops.Tensor]]] = None, feed_dict_creator: Optional[Callable[int, Dict[str, tensorflow.python.framework.ops.Tensor]]] = None, sess: Optional[tensorflow.python.client.session.Session] = None, sparsity_levels: List[int] = (0.0, 0.2, 0.4, 0.6, 0.7, 0.8, 0.85, 0.9, 0.95, 0.99), show_progress: bool = True) → sparseml.optim.sensitivity.PruningLossSensitivityAnalysis

Run a one shot sensitivity analysis for kernel sparsity. It does not retrain, and instead puts the model to eval mode. Moves operation by operation to calculate the sensitivity analysis for each and resets the previously run layers. Subsequent sparsity checks for layers and levels will be much faster.

Note: this should be run once a session has been created and the variables have been created for the model.

Note: the graph should be recreated for later training as this creates extra ops in the graph that should be reused before continuing in the system.

Parameters

• op_vars – the created pruning op vars from ks_loss_sensitivity_op_vars

• loss_tensor – the loss tensor in the model to measure for the sensitivity

• steps_per_measurement – the number of session.run calls to run through for each sparsity level on each layer

• add_ops_creator – a callback to create an op/tens list to be run through the session for each measurement. Called for each measurement

• feed_dict_creator – a callback to create a feed dict to be run through the session for each measurement. Called for each measurement

• sess – the session to use

• sparsity_levels – the sparsity levels to check for each layer to calculate sensitivity

• show_progress – track progress of the runs if True

Returns

the sensitivity results for every op that is prunable

sparseml.tensorflow_v1.optim.sensitivity_pruning.pruning_loss_sens_op_vars(graph: Optional[tensorflow.python.framework.ops.Graph] = None, var_names: Union[List[str], Tuple[str]] = ('re:.*'), mask_type: Union[str, List[int], sparseml.tensorflow_v1.optim.mask_creator_pruning.PruningMaskCreator] = 'unstructured') → List[sparseml.tensorflow_v1.optim.sensitivity_pruning.SparsePruningOpVars]

Edit the graph for to inject pruning ops and vars to allow for a ks loss sensitivity analysis.

Note: this must be run outside of a session for it to take effect.

Parameters

• graph – the graph to inject pruning ops and vars into, if not supplied uses get_default_graph()

• var_names – List of variable names or regex patterns of variables to get the op vars for. Defaults to matching all variables

• 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, or a SparsityMaskCreator object. default is ‘unstructured’

Returns

the created pruning op vars to be used in approx_ks_loss_sensitivity and one_shot_ks_loss_sensitivity

Module contents

Recalibration code for the TensorFlow framework. Handles things like model pruning and increasing activation sparsity.