Join us at TensorFlow World, Oct 28-31. Use code TF20 for 20% off select passes. Register now

tfmot.sparsity.keras.strip_pruning

Strip pruning wrappers from the model.

tfmot.sparsity.keras.strip_pruning(model)

Defined in python/core/sparsity/keras/prune.py.

Once a model has been pruned to required sparsity, this method can be used to restore the original model with the sparse weights.

Only sequential and functional models are supported for now.

Arguments:

  • model: A tf.keras.Model instance with pruned layers.

Returns:

A keras model with pruning wrappers removed.

Raises:

  • ValueError: if the model is not a tf.keras.Model instance.
  • NotImplementedError: if the model is a subclass model.

Usage:

orig_model = tf.keras.Model(inputs, outputs)
pruned_model = prune_low_magnitude(orig_model)
exported_model = strip_pruning(pruned_model)

The exported_model and the orig_model share the same structure.