Make the PruneForLatencyOnXNNPack prune policy more robust

PiperOrigin-RevId: 373696218

Author: tensorflower-gardener

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

@@ -103,7 +103,12 @@ def _get_producers(self, layer):
     return producers
 
   def _get_consumers(self, layer):
-    return [node.outbound_layer for node in layer._outbound_nodes]
+
+    def unpack(layer):
+      return (unpack(layer.layers[0])
+              if isinstance(layer, tf.keras.Sequential) else layer)
+
+    return [unpack(node.outbound_layer) for node in layer._outbound_nodes]
 
   def _lookup_layers(self, source_layers, stop_fn, next_fn):
     """Traverses the model and returns layers satisfying `stop_fn` criteria."""
@@ -127,9 +132,15 @@ def _lookup_layers(self, source_layers, stop_fn, next_fn):
 
   def _start_layer_stop_fn(self, layer):
     """Determines whether the layer starts a subgraph of sparse inference."""
-    return (isinstance(layer, layers.Conv2D) and hasattr(layer, 'kernel') and
-            layer.kernel.shape[:3] == (3, 3, 3) and layer.strides == (2, 2) and
-            layer.padding.lower() == 'valid')
+    if isinstance(layer, layers.Conv2D):
+      producers = self._get_producers(layer)
+      return (hasattr(layer, 'kernel') and
+              layer.kernel.shape[:3] == (3, 3, 3) and
+              layer.strides == (2, 2) and layer.padding.lower() == 'valid' and
+              len(producers) == 1 and
+              isinstance(producers[0], layers.ZeroPadding2D) and
+              producers[0].padding == ((1, 1), (1, 1)))
+    return False
 
   def _end_layer_stop_fn(self, layer):
     """Determines whether the layer ends a subgraph of sparse inference."""
@@ -157,10 +168,36 @@ def _check_layer_support(self, layer):
       # 3x3 stride-2 convolution (no dilation, padding 1 on each side).
       # 5x5 stride-1 convolution (no dilation, padding 2 on each side).
       # 5x5 stride-2 convolution (no dilation, padding 2 on each side).
-      return (layer.depth_multiplier == 1 and layer.dilation_rate == (1, 1) and
-              (layer.kernel_size == (3, 3) or layer.kernel_size == (5, 5)) and
-              ((layer.padding.lower() == 'same' and layer.strides == (1, 1)) or
-               (layer.padding.lower() == 'valid' and layer.strides == (2, 2))))
+      padding = layer.padding.lower()
+      producers = self._get_producers(layer)
+      zero_padding = (
+          producers[0] if len(producers) == 1 and
+          isinstance(producers[0], layers.ZeroPadding2D) else None)
+
+      supported_case_1 = (
+          layer.kernel_size == (3, 3) and layer.strides == (1, 1) and
+          padding == 'same')
+
+      supported_case_2 = (
+          layer.kernel_size == (3, 3) and layer.strides == (2, 2) and
+          padding == 'valid' and zero_padding and
+          zero_padding.padding == ((1, 1), (1, 1)))
+
+      supported_case_3 = (
+          layer.kernel_size == (5, 5) and layer.strides == (1, 1) and
+          padding == 'same')
+
+      supported_case_4 = (
+          layer.kernel_size == (5, 5) and layer.strides == (2, 2) and
+          padding == 'valid' and zero_padding and
+          zero_padding.padding == ((2, 2), (2, 2)))
+
+      supported = (
+          layer.depth_multiplier == 1 and layer.dilation_rate == (1, 1) and
+          (supported_case_1 or supported_case_2 or supported_case_3 or
+           supported_case_4))
+
+      return supported
     elif isinstance(layer, layers.Conv2D):
       # 1x1 convolution (no stride, no dilation, no padding, no groups).
       return (layer.groups == 1 and layer.dilation_rate == (1, 1) and
@@ -200,8 +237,9 @@ def ensure_model_supports_pruning(self, model):
     )
     if not start_layers:
       raise ValueError(('Could not find `Conv2D 3x3` layer with stride 2x2, '
-                        '`input filters == 3` and `VALID` padding in all input '
-                        'branches of the model'))
+                        '`input filters == 3` and `VALID` padding and '
+                        'preceding `ZeroPadding2D` with `padding == 1` in all '
+                        'input branches of the model'))
 
     # Search for the end layer (GlobalAveragePooling with `keepdims = True`)
     # for every output branch (backward).

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

@@ -27,8 +27,10 @@
 
 class PruningPolicyTest(tf.test.TestCase):
   INVALID_TO_PRUNE_START_LAYER_ERROR = (
-      'Could not find `Conv2D 3x3` layer with stride 2x2, `input filters == 3` '
-      'and `VALID` padding in all input branches of the model')
+      'Could not find `Conv2D 3x3` layer with stride 2x2, `input filters == 3`'
+      ' and `VALID` padding and preceding `ZeroPadding2D` with `padding == 1` '
+      'in all input branches of the model'
+  )
 
   INVALID_TO_PRUNE_STOP_LAYER_ERROR = (
       'Could not find a `GlobalAveragePooling2D` layer with `keepdims = True` '
@@ -164,18 +166,18 @@ def testPruneSequentialModelForLatencyOnXNNPackPolicy(self):
     self.assertEqual(self._count_pruned_layers(pruned_model), 1)
 
   def testPruneModelRecursivelyForLatencyOnXNNPackPolicy(self):
-    internal_model = keras.Sequential(
-        [layers.ZeroPadding2D(padding=1, input_shape=(8, 8, 3))])
     original_model = keras.Sequential([
-        internal_model,
+        layers.ZeroPadding2D(padding=1, input_shape=(8, 8, 3)),
         layers.Conv2D(
             filters=4,
             kernel_size=(3, 3),
             strides=(2, 2),
             padding='valid',
         ),
-        layers.Conv2D(filters=8, kernel_size=[1, 1]),
-        layers.Conv2D(filters=16, kernel_size=[1, 1]),
+        keras.Sequential([
+            layers.Conv2D(filters=8, kernel_size=[1, 1]),
+            layers.Conv2D(filters=16, kernel_size=[1, 1]),
+        ]),
         layers.GlobalAveragePooling2D(keepdims=True),
     ])
     pruned_model = prune.prune_low_magnitude(