Fix the pruning policy to handle models after cloning

PiperOrigin-RevId: 375601956

Author: tensorflower-gardener

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

@@ -94,12 +94,12 @@ def allow_pruning(self, layer):
     return isinstance(layer, layers.Conv2D) and layer.kernel_size == (1, 1)
 
   def _get_producers(self, layer):
-    producers = []
+    producers = set()
     for node in layer._inbound_nodes:
       if isinstance(node.inbound_layers, list):
-        producers.extend(node.inbound_layers)
+        producers.update(node.inbound_layers)
       else:
-        producers.append(node.inbound_layers)
+        producers.add(node.inbound_layers)
     return producers
 
   def _get_consumers(self, layer):
@@ -133,7 +133,7 @@ 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."""
     if isinstance(layer, layers.Conv2D):
-      producers = self._get_producers(layer)
+      producers = list(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
@@ -172,7 +172,7 @@ def _check_layer_support(self, layer):
       # 5x5 convolution with `VALID` padding (no dilation, stride-1 or stride-2,
       #   preceding `ZeroPadding2D` layer with padding 2 on each side.
       padding = layer.padding.lower()
-      producers = self._get_producers(layer)
+      producers = list(self._get_producers(layer))
       zero_padding = (
           producers[0] if len(producers) == 1 and
           isinstance(producers[0], layers.ZeroPadding2D) else None)

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

@@ -278,6 +278,29 @@ def testFunctionalModelForLatencyOnXNNPackPolicy(self):
     )
     self.assertEqual(self._count_pruned_layers(pruned_model), 6)
 
+  def testPruneFunctionalModelAfterCloneForLatencyOnXNNPackPolicy(self):
+    i = keras.Input(shape=(8, 8, 3))
+    x = layers.ZeroPadding2D(padding=1)(i)
+    x = layers.Conv2D(
+        filters=16,
+        kernel_size=(3, 3),
+        strides=(2, 2),
+        padding='valid',
+    )(
+        x)
+    x = layers.Conv2D(filters=16, kernel_size=[1, 1])(x)
+    o = layers.GlobalAveragePooling2D(keepdims=True)(x)
+    original_model = keras.Model(inputs=[i], outputs=[o])
+
+    cloned_model = tf.keras.models.clone_model(
+        original_model, clone_function=lambda l: l)
+    pruned_model = prune.prune_low_magnitude(
+        cloned_model,
+        pruning_policy=pruning_policy.PruneForLatencyOnXNNPack(),
+        **self.params,
+    )
+    self.assertEqual(self._count_pruned_layers(pruned_model), 1)
+
 
 if __name__ == '__main__':
   tf.test.main()