Add compat_v1 batchnorm to PruneRegistry

PiperOrigin-RevId: 419516112

Author: rino20

tensorflow_model_optimization/python/core/sparsity/keras/BUILD

@@ -263,6 +263,7 @@ py_strict_test(
     deps = [
         ":prunable_layer",
         ":prune_registry",
+        # absl/testing:parameterized dep1,
         # tensorflow dep1,
     ],
 )

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

@@ -21,6 +21,7 @@
 from tensorflow_model_optimization.python.core.sparsity.keras import prunable_layer
 
 layers = tf.keras.layers
+layers_compat_v1 = tf.compat.v1.keras.layers
 
 
 class PruneRegistry(object):
@@ -94,12 +95,12 @@ class PruneRegistry(object):
       layers.MaxPooling2D: [],
       layers.MaxPooling3D: [],
       layers.MultiHeadAttention: [
-        '_query_dense.kernel',
-        '_key_dense.kernel',
-        '_value_dense.kernel',
-        '_output_dense.kernel'],
+          '_query_dense.kernel', '_key_dense.kernel', '_value_dense.kernel',
+          '_output_dense.kernel'
+      ],
       layers.experimental.preprocessing.Rescaling.__class__: [],
       TensorFlowOpLayer: [],
+      layers_compat_v1.BatchNormalization: [],
   }
 
   _RNN_CELLS_WEIGHTS_MAP = {

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

@@ -14,6 +14,7 @@
 # ==============================================================================
 """Tests for prune registry."""
 
+from absl.testing import parameterized
 import tensorflow as tf
 
 from tensorflow_model_optimization.python.core.sparsity.keras import prunable_layer
@@ -24,102 +25,98 @@
 PruneRegistry = prune_registry.PruneRegistry
 
 
-class PruneRegistryTest(tf.test.TestCase):
-
-  class CustomLayer(layers.Layer):
-    pass
-
-  class CustomLayerFromPrunableLayer(layers.Dense):
-    pass
-
-  class MinimalRNNCell(keras.layers.Layer):
-
-    def __init__(self, units, **kwargs):
-      self.units = units
-      self.state_size = units
-      super(PruneRegistryTest.MinimalRNNCell, self).__init__(**kwargs)
-
-    def build(self, input_shape):
-      self.kernel = self.add_weight(shape=(input_shape[-1], self.units),
-                                    initializer='uniform',
-                                    name='kernel')
-      self.recurrent_kernel = self.add_weight(
-          shape=(self.units, self.units),
-          initializer='uniform',
-          name='recurrent_kernel')
-      self.built = True
-
-    def call(self, inputs, states):
-      prev_output = states[0]
-      h = keras.backend.dot(inputs, self.kernel)
-      output = h + keras.backend.dot(prev_output, self.recurrent_kernel)
-      return output, [output]
-
-  class MinimalRNNCellPrunable(MinimalRNNCell, prunable_layer.PrunableLayer):
-
-    def get_prunable_weights(self):
-      return [self.kernel, self.recurrent_kernel]
-
-  def testSupportsKerasPrunableLayer(self):
-    self.assertTrue(PruneRegistry.supports(layers.Dense(10)))
-
-  def testSupportsKerasPrunableLayerAlias(self):
-    # layers.Conv2D maps to layers.convolutional.Conv2D
-    self.assertTrue(PruneRegistry.supports(layers.Conv2D(10, 5)))
-
-  def testSupportsKerasNonPrunableLayer(self):
-    # Dropout is a layer known to not be prunable.
-    self.assertTrue(PruneRegistry.supports(layers.Dropout(0.5)))
-
-  def testDoesNotSupportKerasUnsupportedLayer(self):
-    # ConvLSTM2D is a built-in keras layer but not supported.
-    self.assertFalse(PruneRegistry.supports(layers.ConvLSTM2D(2, (5, 5))))
-
-  def testSupportsKerasRNNLayers(self):
-    self.assertTrue(PruneRegistry.supports(layers.LSTM(10)))
-    self.assertTrue(PruneRegistry.supports(layers.GRU(10)))
-    self.assertTrue(PruneRegistry.supports(layers.SimpleRNN(10)))
-
-  def testSupportsKerasRNNLayerWithRNNCellsParams(self):
-    self.assertTrue(PruneRegistry.supports(layers.RNN(layers.LSTMCell(10))))
-
-    self.assertTrue(
-        PruneRegistry.supports(
-            layers.RNN([
-                layers.LSTMCell(10),
-                layers.GRUCell(10),
-                keras.experimental.PeepholeLSTMCell(10),
-                layers.SimpleRNNCell(10)
-            ])))
-
-  def testDoesNotSupportKerasRNNLayerUnknownCell(self):
-    self.assertFalse(PruneRegistry.supports(
-        keras.layers.RNN(PruneRegistryTest.MinimalRNNCell(32))))
-
-  def testSupportsKerasRNNLayerPrunableCell(self):
-    self.assertTrue(PruneRegistry.supports(
-        keras.layers.RNN(PruneRegistryTest.MinimalRNNCellPrunable(32))))
-
-  def testDoesNotSupportCustomLayer(self):
-    self.assertFalse(PruneRegistry.supports(PruneRegistryTest.CustomLayer()))
-
-  def testDoesNotSupportCustomLayerInheritedFromPrunableLayer(self):
-    self.assertFalse(
-        PruneRegistry.supports(
-            PruneRegistryTest.CustomLayerFromPrunableLayer(10)))
+class CustomLayer(layers.Layer):
+  pass
+
+
+class CustomLayerFromPrunableLayer(layers.Dense):
+  pass
+
+
+class MinimalRNNCell(keras.layers.Layer):
+
+  def __init__(self, units, **kwargs):
+    self.units = units
+    self.state_size = units
+    super(MinimalRNNCell, self).__init__(**kwargs)
+
+  def build(self, input_shape):
+    self.kernel = self.add_weight(
+        shape=(input_shape[-1], self.units),
+        initializer='uniform',
+        name='kernel')
+    self.recurrent_kernel = self.add_weight(
+        shape=(self.units, self.units),
+        initializer='uniform',
+        name='recurrent_kernel')
+    self.built = True
+
+  def call(self, inputs, states):
+    prev_output = states[0]
+    h = keras.backend.dot(inputs, self.kernel)
+    output = h + keras.backend.dot(prev_output, self.recurrent_kernel)
+    return output, [output]
+
+
+class MinimalRNNCellPrunable(MinimalRNNCell, prunable_layer.PrunableLayer):
+
+  def get_prunable_weights(self):
+    return [self.kernel, self.recurrent_kernel]
+
+
+class PruneRegistryTest(tf.test.TestCase, parameterized.TestCase):
+
+  _PRUNE_REGISTRY_SUPPORTED_LAYERS = [
+      # Supports basic Keras layers even though it is not prunbale.
+      layers.Dense(10),
+      layers.Conv2D(10, 5),
+      layers.Dropout(0.5),
+      # Supports specific layers from experimental or compat_v1.
+      tf.keras.layers.experimental.preprocessing.Rescaling,
+      tf.compat.v1.keras.layers.BatchNormalization(),
+      # Supports Keras RNN Layers with prunable cells.
+      layers.LSTM(10),
+      layers.GRU(10),
+      layers.SimpleRNN(10),
+      layers.RNN(layers.LSTMCell(10)),
+      layers.RNN([
+          layers.LSTMCell(10),
+          layers.GRUCell(10),
+          keras.experimental.PeepholeLSTMCell(10),
+          layers.SimpleRNNCell(10)
+      ]),
+      keras.layers.RNN(MinimalRNNCellPrunable(32)),
+  ]
+
+  @parameterized.parameters(_PRUNE_REGISTRY_SUPPORTED_LAYERS)
+  def testSupportsLayer(self, layer):
+    self.assertTrue(PruneRegistry.supports(layer))
+
+  _PRUNE_REGISTRY_UNSUPPORTED_LAYERS = [
+      # Not support a few built-in keras layers.
+      layers.ConvLSTM2D(2, (5, 5)),
+      # Not support RNN layers with unknown cell
+      keras.layers.RNN(MinimalRNNCell(32)),
+      # Not support Custom layers, even though inherited from prunable layer.
+      CustomLayer(),
+      CustomLayerFromPrunableLayer(10),
+  ]
+
+  @parameterized.parameters(_PRUNE_REGISTRY_UNSUPPORTED_LAYERS)
+  def testDoesNotSupportLayer(self, layer):
+    self.assertFalse(PruneRegistry.supports(layer))
 
   def testMakePrunableRaisesErrorForKerasUnsupportedLayer(self):
     with self.assertRaises(ValueError):
       PruneRegistry.make_prunable(layers.ConvLSTM2D(2, (5, 5)))
 
   def testMakePrunableRaisesErrorForCustomLayer(self):
     with self.assertRaises(ValueError):
-      PruneRegistry.make_prunable(PruneRegistryTest.CustomLayer())
+      PruneRegistry.make_prunable(CustomLayer())
 
   def testMakePrunableRaisesErrorForCustomLayerInheritedFromPrunableLayer(self):
     with self.assertRaises(ValueError):
-      PruneRegistry.make_prunable(
-          PruneRegistryTest.CustomLayerFromPrunableLayer(10))
+      PruneRegistry.make_prunable(CustomLayerFromPrunableLayer(10))
 
   def testMakePrunableWorksOnKerasPrunableLayer(self):
     layer = layers.Dense(10)
@@ -171,7 +168,7 @@ def testMakePrunableWorksOnKerasRNNLayerWithRNNCellsParams(self):
 
   def testMakePrunableWorksOnKerasRNNLayerWithPrunableCell(self):
     cell1 = layers.LSTMCell(10)
-    cell2 = PruneRegistryTest.MinimalRNNCellPrunable(5)
+    cell2 = MinimalRNNCellPrunable(5)
     layer = layers.RNN([cell1, cell2])
     with self.assertRaises(AttributeError):
       layer.get_prunable_weights()
@@ -187,12 +184,9 @@ def testMakePrunableWorksOnKerasRNNLayerWithPrunableCell(self):
 
   def testMakePrunableRaisesErrorOnRNNLayersUnsupportedCell(self):
     with self.assertRaises(ValueError):
-      PruneRegistry.make_prunable(layers.RNN(
-          [layers.LSTMCell(10), PruneRegistryTest.MinimalRNNCell(5)]))
-
-  def testRescalingLayer(self):
-    self.assertTrue(PruneRegistry.supports(
-        tf.keras.layers.experimental.preprocessing.Rescaling))
+      PruneRegistry.make_prunable(
+          layers.RNN([layers.LSTMCell(10),
+                      MinimalRNNCell(5)]))
 
 
 if __name__ == '__main__':