Ceiling the remaining weights to ensure there is at least 1 connection.

Fixed #215

PiperOrigin-RevId: 338257075

Author: evcu

RELEASE.md

@@ -62,6 +62,7 @@ Tested against TensorFlow nightly, and Python 3.
 Keras pruning API:
 
 Tested against TensorFlow 1.14.0, 2.0.0, and nightly, and Python 3.
+Pruning now doesn't remove the last remaining connection. So extreme sparsities like 0.999.. would remove all connections but one.
 
 
 # TensorFlow Model Optimization 0.4.0

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

@@ -94,12 +94,12 @@ def _get_params_for_layer(layer_type):
         # TODO(tf-mot): fix for Conv2DTranspose on some form of eager,
         # with or without functions. The weights become nan (though the
         # mask seems fine still).
-        #layers.Conv2DTranspose: ([2, (3, 3)], (7, 6, 3)),
+        # layers.Conv2DTranspose: ([2, (3, 3)], (7, 6, 3)),
         layers.Conv3D: ([2, (3, 3, 3)], (5, 7, 6, 3)),
         # TODO(tf-mot): fix for Conv3DTranspose on some form of eager,
         # with or without functions. The weights become nan (though the
         # mask seems fine still).
-        #layers.Conv3DTranspose: ([2, (3, 3, 3)], (5, 7, 6, 3)),
+        # layers.Conv3DTranspose: ([2, (3, 3, 3)], (5, 7, 6, 3)),
         layers.SeparableConv1D: ([4, 3], (3, 6)),
         layers.SeparableConv2D: ([4, (2, 2)], (4, 6, 1)),
         layers.Dense: ([4], (6,)),
@@ -193,17 +193,19 @@ def testPrunesZeroSparsity_IsNoOp(self):
     self._assert_weights_different_objects(model, pruned_model)
     self._assert_weights_same_values(model, pruned_model)
 
-  # TODO(tfmot): https://github.com/tensorflow/model-optimization/issues/215
-  def testPruneWithHighSparsity_Fails(self):
+  def testPruneWithHighSparsity(self):
     params = self.params
     params['pruning_schedule'] = pruning_schedule.ConstantSparsity(
         target_sparsity=0.99, begin_step=0, frequency=1)
 
     model = prune.prune_low_magnitude(
         keras_test_utils.build_simple_dense_model(), **params)
-
-    with self.assertRaises(tf.errors.InvalidArgumentError):
-      self._train_model(model, epochs=1)
+    self._train_model(model, epochs=1)
+    for layer in model.layers:
+      if isinstance(layer, pruning_wrapper.PruneLowMagnitude):
+        for weight in layer.layer.get_prunable_weights():
+          self.assertEqual(
+              1, np.count_nonzero(tf.keras.backend.get_value(weight)))
 
   ###################################################################
   # Tests for training with pruning with pretrained models or weights.

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

@@ -85,9 +85,12 @@ def _update_mask(self, weights):
     with tf.name_scope('pruning_ops'):
       abs_weights = tf.math.abs(weights)
       k = tf.dtypes.cast(
-          tf.math.round(
-              tf.dtypes.cast(tf.size(abs_weights), tf.float32) *
-              (1 - sparsity)), tf.int32)
+          tf.math.maximum(
+              tf.math.round(
+                  tf.dtypes.cast(tf.size(abs_weights), tf.float32) *
+                  (1 - sparsity)),
+              1),
+          tf.int32)
       # Sort the entire array
       values, _ = tf.math.top_k(
           tf.reshape(abs_weights, [-1]), k=tf.size(abs_weights))

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

@@ -71,6 +71,37 @@ def training_step_fn():
 
     compat.initialize_variables(self)
 
+  def testExtremelySparseMask(self):
+    weight = tf.Variable(np.linspace(1.0, 100.0, 100), name="weights")
+    weight_dtype = weight.dtype.base_dtype
+    mask = tf.Variable(
+        tf.ones(weight.get_shape(), dtype=weight_dtype),
+        name="mask",
+        dtype=weight_dtype)
+    threshold = tf.Variable(
+        tf.zeros([], dtype=weight_dtype), name="threshold", dtype=weight_dtype)
+    self.initialize()
+
+    extreme_sparsity = pruning_schedule.ConstantSparsity(0.9999, 0, 100, 1)
+    p = pruning_impl.Pruning(
+        pruning_vars=[(weight, mask, threshold)],
+        training_step_fn=self.training_step_fn,
+        pruning_schedule=extreme_sparsity,
+        block_size=self.block_size,
+        block_pooling_type=self.block_pooling_type)
+
+    mask_before_pruning = K.get_value(mask)
+    self.assertAllEqual(np.count_nonzero(mask_before_pruning), 100)
+
+    if tf.executing_eagerly():
+      p.conditional_mask_update()
+    else:
+      K.get_session().run(p.conditional_mask_update())
+
+    # We should always have a single connection remaining.
+    mask_after_pruning = K.get_value(mask)
+    self.assertAllEqual(np.count_nonzero(mask_after_pruning), 1)
+
   def testUpdateSingleMask(self):
     weight = tf.Variable(np.linspace(1.0, 100.0, 100), name="weights")
     weight_dtype = weight.dtype.base_dtype
@@ -154,8 +185,7 @@ def testBlockMaskingMax(self):
     block_size = (2, 2)
     block_pooling_type = "MAX"
     weight = tf.constant([[0.1, 0.0, 0.2, 0.0], [0.0, -0.1, 0.0, -0.2],
-                                   [0.3, 0.0, 0.4, 0.0], [0.0, -0.3, 0.0,
-                                                          -0.4]])
+                          [0.3, 0.0, 0.4, 0.0], [0.0, -0.3, 0.0, -0.4]])
     expected_mask = [[0.0, 0.0, 0.0, 0.0], [0.0, 0.0, 0.0, 0.0],
                      [1., 1., 1., 1.], [1., 1., 1., 1.]]
 
@@ -167,8 +197,7 @@ def testBlockMaskingWithHigherDimensionsRaisesError(self):
     block_pooling_type = "AVG"
     # Weights as in testBlockMasking, but with one extra dimension.
     weight = tf.constant([[[0.1, 0.1, 0.2, 0.2], [0.1, 0.1, 0.2, 0.2],
-                                    [0.3, 0.3, 0.4, 0.4], [0.3, 0.3, 0.4,
-                                                           0.4]]])
+                           [0.3, 0.3, 0.4, 0.4], [0.3, 0.3, 0.4, 0.4]]])
     expected_mask = [[[0.0, 0.0, 0.0, 0.0], [0.0, 0.0, 0.0, 0.0],
                       [1., 1., 1., 1.], [1., 1., 1., 1.]]]