sparseml.keras.optim package

Submodules

sparseml.keras.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.keras.optim.manager.ScheduledModifierManager(modifiers: List[sparseml.keras.optim.modifier.ScheduledModifier])[source]

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

The base modifier manager, handles managing multiple ScheduledModifier.

finalize(model: tensorflow.python.keras.engine.training.Model)[source]

Remove extra information related to the modifier from the model that is not necessary for exporting

Parameters

model – a Keras model

Returns

a new Keras model

static from_yaml(file_path: Union[str, sparsezoo.objects.recipe.Recipe], add_modifiers: Optional[List[sparseml.keras.optim.modifier.Modifier]] = None, **recipe_variables)[source]

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)

Returns

ScheduledModifierManager() created from the recipe file

modify(model: tensorflow.python.keras.engine.training.Model, optimizer: tensorflow.python.keras.optimizer_v2.optimizer_v2.OptimizerV2, steps_per_epoch: int, loggers: Optional[Union[sparseml.keras.utils.logger.KerasLogger, List[sparseml.keras.utils.logger.KerasLogger]]] = None, input_tensors: Optional[tensorflow.python.framework.ops.Tensor] = None)[source]

Modify the model and optimizer based on the requirements of modifiers

Parameters
  • model – model to modify

  • optimizer – optimizer to modify

  • steps_per_epoch – number of steps per epoch

  • loggers – list of loggers

  • input_tensors – optional input tensor

Returns

model, optimizer, callbacks

sparseml.keras.optim.mask_pruning module

class sparseml.keras.optim.mask_pruning.MaskedLayer(*args, **kwargs)[source]

Bases: tensorflow.python.keras.layers.wrappers.Wrapper

Masked layer is a layer wrapping around another layer with a mask; the mask however is shared if the enclosed layer is again of MaskedLayer type

Parameters
  • layer – either a MaskedLayer or a keras layer

  • pruning_scheduler – a pruning scheduler

  • mask_creator – a mask creator

  • kwargs – optional params for keras layer constructor, e.g. layer name

build(input_shape)[source]

Creates the variables of the layer (optional, for subclass implementers).

This is a method that implementers of subclasses of Layer or Model can override if they need a state-creation step in-between layer instantiation and layer call.

This is typically used to create the weights of Layer subclasses.

Parameters

input_shape – Instance of TensorShape, or list of instances of TensorShape if the layer expects a list of inputs (one instance per input).

call(inputs: tensorflow.python.framework.ops.Tensor, training=None)[source]

Forward function for calling layer instance as function

compute_output_shape(input_shape)[source]

Computes the output shape of the layer.

If the layer has not been built, this method will call build on the layer. This assumes that the layer will later be used with inputs that match the input shape provided here.

Parameters

input_shape – Shape tuple (tuple of integers) or list of shape tuples (one per output tensor of the layer). Shape tuples can include None for free dimensions, instead of an integer.

Returns

An input shape tuple.

classmethod from_config(config)[source]

Creates a layer from its config.

This method is the reverse of get_config, capable of instantiating the same layer from the config dictionary. It does not handle layer connectivity (handled by Network), nor weights (handled by set_weights).

Parameters

config – A Python dictionary, typically the output of get_config.

Returns

A layer instance.

get_config()[source]

Get layer config Serialization and deserialization should be done using keras.serialize/deserialize, which create and retrieve the “class_name” field automatically. The resulting config below therefore does not contain the field.

property global_step
property mask_updater
property masked_layer
property masks
property pruned_layer
property pruning_vars
class sparseml.keras.optim.mask_pruning.PruningScheduler[source]

Bases: abc.ABC

Abstract pruning scheduler

classmethod deserialize(config)[source]

Deserialize a pruning scheduler from config returned by scheduler’s get_config method

Parameters

config – a pruning scheduler’s config

Returns

a pruning scheduler instance

abstract get_config()[source]
abstract should_prune(step: int)bool[source]

Check if the given step is a right time for pruning

Parameters

step – training step

Returns

True if pruning should take place; False otherwise

abstract target_sparsity(step: int, **kwargs)float[source]

Compute the target sparsity at the given step

Parameters
  • step – training step

  • kwargs – optional keyword params that a specific scheduler might need

Returns

target sparsity

sparseml.keras.optim.mask_pruning.remove_pruning_masks(model: tensorflow.python.keras.engine.training.Model)[source]

Remove pruning masks from a model that was pruned using the MaskedLayer logic :param model: a model that was pruned using MaskedLayer :return: the original model with pruned weights

sparseml.keras.optim.mask_pruning_creator module

Classes for defining sparsity masks based on model parameters.

class sparseml.keras.optim.mask_pruning_creator.BlockPruningMaskCreator(block_shape: List[int], grouping_op_name: str = 'mean')[source]

Bases: sparseml.keras.optim.mask_pruning_creator.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[source]
Parameters

tensor – The tensor to transform

Returns

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

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

Bases: sparseml.keras.optim.mask_pruning_creator.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[source]
Parameters

tensor – The tensor to transform

Returns

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

class sparseml.keras.optim.mask_pruning_creator.GroupedPruningMaskCreator[source]

Bases: sparseml.keras.optim.mask_pruning_creator.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[source]
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[source]
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][source]
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[source]
Parameters

tensor – The tensor to reduce in groups

Returns

The grouped tensor

class sparseml.keras.optim.mask_pruning_creator.PruningMaskCreator[source]

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[source]
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][source]
Parameters

tensor – A tensor of a model layer’s weights

Returns

Tensor initializer function for this sparsity mask

class sparseml.keras.optim.mask_pruning_creator.UnstructuredPruningMaskCreator[source]

Bases: sparseml.keras.optim.mask_pruning_creator.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[source]
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][source]
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.keras.optim.mask_pruning_creator.load_mask_creator(obj: Union[str, Iterable[int]])sparseml.keras.optim.mask_pruning_creator.PruningMaskCreator[source]
Parameters

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

Returns

SparsityMaskCreator object created from obj

sparseml.keras.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.keras.optim.modifier.KerasModifierYAML[source]

Bases: sparseml.optim.modifier.ModifierYAML

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

class sparseml.keras.optim.modifier.Modifier(log_types: Optional[Union[str, List[str]]] = None, **kwargs)[source]

Bases: sparseml.optim.modifier.BaseModifier

Base modifier class that all Keras modifiers should derive themselves from. Handles setting up the expected contracts for modifying model and optimizer

Modifiers are expected to implement the following functions for Keras:
- modify - modify model and optimizer
Parameters
  • log_types – the loggers that can be used by the modifier instance

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

finalize(model: tensorflow.python.keras.engine.training.Model)[source]

Remove extra information related to the modifier from the model that is not necessary for exporting

Parameters

model – a Keras model

Returns

a new Keras model

static load_list(yaml_str: str)[source]
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)[source]
Parameters

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

Returns

the loaded modifier object

modify(model, optimizer, steps_per_epoch: int, loggers: Optional[Union[sparseml.keras.utils.logger.KerasLogger, List[sparseml.keras.utils.logger.KerasLogger]]] = None, input_tensors: Optional[tensorflow.python.framework.ops.Tensor] = None)[source]

Modify model, optimizer based on the logic of the modifier. Return the modified model, optimizer and a list of callbacks (e.g., to enhance training process)

Parameters
  • model – model to modify

  • optimizer – optimizer to modify

  • steps_per_epoch – number of steps per epoch

  • input_tensors – optional input tensor

Returns

model, optimizer, callbacks

class sparseml.keras.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)[source]

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

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

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.keras.optim.modifier.ScheduledModifier(log_types: Optional[Union[str, List[str]]] = None, 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)[source]

Bases: sparseml.keras.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 the following functions for Keras:
- modify - modify model and optimizer
Parameters
  • log_types – the loggers that can be used by the modifier instance

  • 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 less than, 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, after_optim: bool)Tuple[int, int][source]

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.keras.optim.modifier.ScheduledUpdateModifier(log_types: Optional[Union[str, List[str]]] = None, 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)[source]

Bases: sparseml.keras.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 the following functions for Keras:
- modify - modify model and optimizer
Parameters
  • log_types – the loggers that can be used by the modifier instance

  • 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 less than, 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[source]

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)

sparseml.keras.optim.modifier_epoch module

Contains code for epoch modifiers in Keras

class sparseml.keras.optim.modifier_epoch.EpochRangeModifier(start_epoch: float, end_epoch: float)[source]

Bases: sparseml.sparsification.modifier_epoch.EpochRangeModifier, sparseml.keras.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.keras.optim.modifier_lr module

Learning rate modifiers for Keras models

class sparseml.keras.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, log_types: Union[str, List[str]] = '__ALL__')[source]

Bases: sparseml.sparsification.modifier_lr.LearningRateModifier, sparseml.keras.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
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__

modify(model, optimizer, steps_per_epoch: int, loggers: Optional[Union[sparseml.keras.utils.logger.KerasLogger, List[sparseml.keras.utils.logger.KerasLogger]]] = None, input_tensors: Optional[tensorflow.python.framework.ops.Tensor] = None)[source]

Modify model and optimizer, and provide callbacks to process the model

Parameters
  • model – a model to be modified with prunable layers wrapped by masks

  • optimizer – an optimizer to be modified

  • steps_per_epoch – number of steps per epoch

  • loggers – list of loggers

  • input_tensors – optional input tensors

Returns

modified model, optimizer and callbacks

class sparseml.keras.optim.modifier_lr.SetLearningRateModifier(learning_rate: float, start_epoch: float = - 1, end_epoch: float = - 1, log_types: Union[str, List[str]] = '__ALL__')[source]

Bases: sparseml.sparsification.modifier_lr.SetLearningRateModifier, sparseml.keras.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
log_types: __ALL__
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

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

modify(model, optimizer, steps_per_epoch: int, loggers: Optional[Union[sparseml.keras.utils.logger.KerasLogger, List[sparseml.keras.utils.logger.KerasLogger]]] = None, input_tensors: Optional[tensorflow.python.framework.ops.Tensor] = None)[source]

Modify model and optimizer, and provide callbacks to process the model

Parameters
  • model – a model to be modified with prunable layers wrapped by masks

  • optimizer – an optimizer to be modified

  • steps_per_epoch – number of steps per epoch

  • loggers – list of loggers

  • input_tensors – optional input tensors

Returns

modified model, optimizer and callbacks

sparseml.keras.optim.modifier_params module

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

class sparseml.keras.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)[source]

Bases: sparseml.sparsification.modifier_params.TrainableParamsModifier, sparseml.keras.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: [“conv2d_1/kernel:0”, “conv2d_5/kernel:0”]
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

property layer_names
modify(model, optimizer, steps_per_epoch: int, input_tensors: Optional[tensorflow.python.framework.ops.Tensor] = None)[source]

Modify model, optimizer based on the logic of the modifier. Return the modified model, optimizer and a list of callbacks (e.g., to enhance training process)

Parameters
  • model – model to modify

  • optimizer – optimizer to modify

  • steps_per_epoch – number of steps per epoch

  • input_tensors – optional input tensor

Returns

model, optimizer, callbacks

validate()[source]

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

sparseml.keras.optim.modifier_pruning module

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

class sparseml.keras.optim.modifier_pruning.ConstantPruningModifier(params: Union[str, List[str]], start_epoch: float = - 1, end_epoch: float = - 1, log_types: Union[str, List[str]] = '__ALL__')[source]

Bases: sparseml.sparsification.modifier_pruning.ConstantPruningModifier, sparseml.keras.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
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__

finalize(model: tensorflow.python.keras.engine.training.Model)[source]

Remove extra information related to the modifier from the model that is not necessary for exporting

Parameters

model – a Keras model

Returns

a new Keras model

is_pruning_step(step: int, steps_per_epoch, tensor=None)[source]
property layer_names
modify(model, optimizer, steps_per_epoch: int, loggers: Optional[Union[sparseml.keras.utils.logger.KerasLogger, List[sparseml.keras.utils.logger.KerasLogger]]] = None, input_tensors: Optional[tensorflow.python.framework.ops.Tensor] = None)[source]

Modify model and optimizer

Parameters
  • model – a model to be modified

  • optimizer – an optimizer to be modified

  • steps_per_epoch – number of steps per epoch

  • loggers – list of loggers

  • input_tensors – optional input tensors

Returns

modified model, optimizer and callbacks

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

Bases: sparseml.sparsification.modifier_pruning.GMPruningModifier, sparseml.keras.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
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

finalize(model: tensorflow.python.keras.engine.training.Model)[source]

Remove extra information related to the modifier from the model that is not necessary for exporting

Parameters

model – a Keras model

Returns

a new Keras model

property layer_names
modify(model, optimizer, steps_per_epoch: int, loggers: Optional[Union[sparseml.keras.utils.logger.KerasLogger, List[sparseml.keras.utils.logger.KerasLogger]]] = None, input_tensors: Optional[tensorflow.python.framework.ops.Tensor] = None)[source]

Modify model and optimizer, and provide callbacks to process the model

Parameters
  • model – a model to be modified with prunable layers wrapped by masks

  • optimizer – an optimizer to be modified; in this case, no change to it

  • steps_per_epoch – number of steps per epoch

  • loggers – list of loggers

  • input_tensors – optional input tensors

Returns

modified model, optimizer and callbacks

property prunable_layers
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.keras.optim.utils module

sparseml.keras.optim.utils.get_layer_name_from_param(param: str)[source]

Module contents

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