Module: tfmot.sparsity.keras

View source on GitHub

Module containing sparsity code built on Keras abstractions.


compat module: Functions for TF 1.X and 2.X compatibility.

pruning_sched module: Pruning Schedule classes to control pruning rate during training.

pruning_wrapper module: A Keras wrapper to add pruning related variables to a layer.


class ConstantSparsity: Pruning schedule with constant sparsity(%) throughout training.

class PolynomialDecay: Pruning Schedule with a PolynomialDecay function.

class PrunableLayer: Abstract Base Class for making your own keras layer prunable.

class PruningSchedule: Specifies when to prune layer and the sparsity(%) at each training step.

class PruningSummaries: A Keras callback for adding pruning summaries to tensorboard.

class UpdatePruningStep: Keras callback which updates pruning wrappers with the optimizer step.


custom_object_scope(...): Provides a scope that changes to _GLOBAL_CUSTOM_OBJECTS cannot escape.

prune_low_magnitude(...): Modify a tf.keras layer or model to be pruned during training.

prune_scope(...): Provides a scope in which Pruned layers and models can be deserialized.

strip_pruning(...): Strip pruning wrappers from the model.