1. The layer is in PruneRegistry (most keras built-in layers).

2. The layer is a `PrunableLayer` instance.
3. Newly added. The layer has a customer defined `get_prunable_weights` method.

PiperOrigin-RevId: 355510603

Author: tensorflower-gardener

tensorflow_model_optimization/python/core/sparsity/keras/pruning_wrapper.py

@@ -67,8 +67,10 @@ class PruneLowMagnitude(Wrapper):
   Custom keras layers:
   The pruning wrapper can also be applied to a user-defined keras layer.
   Such a layer may contain one or more weight tensors that may be pruned.
-  To apply pruning wrapper to such layers, set prunable_weight_names to mark
-  the weight tensors for pruning.
+  To apply pruning wrapper to such layers, the layer should be a `PrunableLayer`
+  instance or, more directly, user should define a `get_prunable_weights` method
+  for the layer (Check the pruning_wrapper_test.CustomLayerPrunable for more
+  details about how to define a user-defined prunable layer).
 
   Sparsity function:
   The target sparsity for the weight tensors are set through the
@@ -141,7 +143,8 @@ def __init__(self,
     kwargs.update({'name': '{}_{}'.format(
         generic_utils.to_snake_case(self.__class__.__name__), layer.name)})
 
-    if isinstance(layer, prunable_layer.PrunableLayer):
+    if isinstance(layer, prunable_layer.PrunableLayer) or hasattr(
+        layer, 'get_prunable_weights'):
       # Custom layer in client code which supports pruning.
       super(PruneLowMagnitude, self).__init__(layer, **kwargs)
     elif prune_registry.PruneRegistry.supports(layer):
@@ -151,8 +154,10 @@ def __init__(self,
     else:
       raise ValueError(
           'Please initialize `Prune` with a supported layer. Layers should '
-          'either be a `PrunableLayer` instance, or should be supported by the '
-          'PruneRegistry. You passed: {input}'.format(input=layer.__class__))
+          'either be supported by the PruneRegistry (built-in keras layers) or '
+          'should be a `PrunableLayer` instance, or should has a customer '
+          'defined `get_prunable_weights` method. You passed: '
+          '{input}'.format(input=layer.__class__))
 
     self._track_trackable(layer, name='layer')
 

tensorflow_model_optimization/python/core/sparsity/keras/pruning_wrapper_test.py

@@ -28,6 +28,35 @@
 Prune = pruning_wrapper.PruneLowMagnitude
 
 
+class CustomLayer(keras.layers.Layer):
+  """A custom layer which is not prunable."""
+
+  def __init__(self, input_dim=16, output_dim=32):
+    super(CustomLayer, self).__init__()
+    self.weight = self.add_weight(
+        shape=(input_dim, output_dim),
+        initializer='random_normal',
+        trainable=True)
+    self.bias = self.add_weight(
+        shape=(output_dim),
+        initializer='zeros',
+        trainable=True)
+
+  def call(self, inputs):
+    return tf.matmul(inputs, self.weight) + self.bias
+
+
+class CustomLayerPrunable(CustomLayer):
+  """A prunable custom layer.
+
+  The layer is same with the CustomLayer except it has a 'get_prunable_weights'
+  attribute.
+  """
+
+  def get_prunable_weights(self):
+    return [self.weight, self.bias]
+
+
 class PruningWrapperTest(tf.test.TestCase):
 
   def setUp(self):
@@ -100,6 +129,19 @@ def testPruneModel(self):
                 'PruneLowMagnitude': pruning_wrapper.PruneLowMagnitude
             }).get_config())
 
+  def testCustomLayerNonPrunable(self):
+    layer = CustomLayer(input_dim=16, output_dim=32)
+    inputs = keras.layers.Input(shape=(16))
+    _ = layer(inputs)
+    with self.assertRaises(ValueError):
+      pruning_wrapper.PruneLowMagnitude(layer, block_pooling_type='MAX')
+
+  def testCustomLayerPrunable(self):
+    layer = CustomLayerPrunable(input_dim=16, output_dim=32)
+    inputs = keras.layers.Input(shape=(16))
+    _ = layer(inputs)
+    pruning_wrapper.PruneLowMagnitude(layer, block_pooling_type='MAX')
+
 
 if __name__ == '__main__':
   tf.test.main()