Skip to content

Testing approach for subclassed models for layer-wise technique. #554

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
wants to merge 1 commit into from
Closed

Testing approach for subclassed models for layer-wise technique. #554

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions tensorflow_model_optimization/python/core/sparsity/keras/BUILD
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -194,6 +194,7 @@ py_test(
# absl/testing:parameterized dep1,
# numpy dep1,
# tensorflow dep1,
"//tensorflow_model_optimization/python/core/keras:compat",
"//tensorflow_model_optimization/python/core/keras:test_utils",
],
)
Expand Down
31 changes: 20 additions & 11 deletions tensorflow_model_optimization/python/core/sparsity/keras/prune.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -194,22 +194,25 @@ def _add_pruning_wrapper(layer):
'an object of type: {input}.'.format(input=to_prune.__class__.__name__))


def strip_pruning(model):
"""Strip pruning wrappers from the model.
def strip_pruning(to_strip):
"""Strip pruning wrappers from the model or layer.

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

Only sequential and functional models are supported for now.

Arguments:
model: A `tf.keras.Model` instance with pruned layers.
to_strip: A `tf.keras.Model` instance with pruned layers or a
`tf.keras.layers.Layer` instance.

Returns:
A keras model with pruning wrappers removed.
A keras model or layer with pruning wrappers removed.

Raises:
ValueError: if the model is not a `tf.keras.Model` instance.
ValueError: if the model is not a `tf.keras.Model` or
`tf.keras.layers.Layer` instance.
NotImplementedError: if the model is a subclass model.

Usage:
Expand All @@ -222,9 +225,11 @@ def strip_pruning(model):
The exported_model and the orig_model share the same structure.
"""

if not isinstance(model, keras.Model):
if not isinstance(to_strip, keras.Model) and not isinstance(
to_strip, keras.layers.Layer):
raise ValueError(
'Expected model to be a `tf.keras.Model` instance but got: ', model)
'Expected `to_strip` to be a `tf.keras.Model` or `tf.keras.layers.Layer` instance but got: ',
to_strip)

def _strip_pruning_wrapper(layer):
if isinstance(layer, tf.keras.Model):
Expand All @@ -241,5 +246,9 @@ def _strip_pruning_wrapper(layer):
return layer.layer
return layer

return keras.models.clone_model(
model, input_tensors=None, clone_function=_strip_pruning_wrapper)
if isinstance(to_strip, keras.Model):
return keras.models.clone_model(
to_strip, input_tensors=None, clone_function=_strip_pruning_wrapper)

if isinstance(to_strip, keras.layers.Layer):
return _strip_pruning_wrapper(to_strip)
Loading