Save mask and threshold to model checkpoint.

PiperOrigin-RevId: 246361345

Author: liyunlu0618

tensorflow_model_optimization/python/core/sparsity/keras/BUILD

@@ -97,6 +97,7 @@ py_library(
         # python:control_flow_ops tensorflow dep2,
         # python:dtypes tensorflow dep2,
         # python:framework_ops tensorflow dep2,
+        # python:variables tensorflow dep2,
         # python/keras:backend tensorflow dep2,
         # python/keras:engine tensorflow dep2,
         # python/keras:generic_utils tensorflow dep2,

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

@@ -259,6 +259,49 @@ def testPruneStopAndRestartOnModel(self, save_restore_fn):
 
     self._check_strip_pruning_matches_original(model, 0.6)
 
+  @parameterized.parameters(test_utils.save_restore_fns())
+  def testPruneWithPolynomialDecayPreservesSparsity(self, save_restore_fn):
+    params = {
+        'pruning_schedule': pruning_schedule.PolynomialDecay(
+            0.2, 0.6, 0, 1, 3, 1),
+        'block_size': (1, 1),
+        'block_pooling_type': 'AVG'
+    }
+    model = prune.prune_low_magnitude(
+        test_utils.build_simple_dense_model(), **params)
+    model.compile(
+        loss='categorical_crossentropy', optimizer='sgd', metrics=['accuracy'])
+    # Model hasn't been trained yet. Sparsity 0.0
+    test_utils.assert_model_sparsity(self, 0.0, model)
+
+    model.fit(
+        np.random.rand(20, 10),
+        keras.utils.to_categorical(np.random.randint(5, size=(20, 1)), 5),
+        batch_size=20,
+        callbacks=[pruning_callbacks.UpdatePruningStep()])
+    # Training has run only 1 step. Sparsity 0.2 (initial_sparsity)
+    test_utils.assert_model_sparsity(self, 0.2, model)
+
+    model.fit(
+        np.random.rand(20, 10),
+        keras.utils.to_categorical(np.random.randint(5, size=(20, 1)), 5),
+        batch_size=20,
+        callbacks=[pruning_callbacks.UpdatePruningStep()])
+    # Training has run 2 steps. Sparsity 0.6 (final_sparsity)
+    test_utils.assert_model_sparsity(self, 0.6, model)
+
+    model = save_restore_fn(model)
+    model.fit(
+        np.random.rand(20, 10),
+        keras.utils.to_categorical(np.random.randint(5, size=(20, 1)), 5),
+        batch_size=20,
+        epochs=2,
+        callbacks=[pruning_callbacks.UpdatePruningStep()])
+    # Training has run all 4 steps. Sparsity 0.6 (final_sparsity)
+    test_utils.assert_model_sparsity(self, 0.6, model)
+
+    self._check_strip_pruning_matches_original(model, 0.6)
+
   def testPrunesPreviouslyUnprunedModel(self):
     model = test_utils.build_simple_dense_model()
     model.compile(

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

@@ -31,6 +31,7 @@
 from tensorflow.python.keras.utils import tf_utils
 from tensorflow.python.ops import check_ops
 from tensorflow.python.ops import control_flow_ops
+from tensorflow.python.ops import variables as tf_variables
 from tensorflow_model_optimization.python.core.sparsity.keras import prunable_layer
 from tensorflow_model_optimization.python.core.sparsity.keras import prune_registry
 from tensorflow_model_optimization.python.core.sparsity.keras import pruning_impl
@@ -175,8 +176,20 @@ def build(self, input_shape):
 
     # For each of the prunable weights, add mask and threshold variables
     for weight in self.prunable_weights:
-      mask = K.variable(np.ones(weight.shape), dtype=weight.dtype)
-      threshold = K.variable(0, dtype=weight.dtype)
+      mask = self.add_variable(
+          'mask',
+          shape=weight.shape,
+          initializer=initializers.get('ones'),
+          dtype=weight.dtype,
+          trainable=False,
+          aggregation=tf_variables.VariableAggregation.MEAN)
+      threshold = self.add_variable(
+          'threshold',
+          shape=[],
+          initializer=initializers.get('zeros'),
+          dtype=weight.dtype,
+          trainable=False,
+          aggregation=tf_variables.VariableAggregation.MEAN)
 
       weight_vars.append(weight)
       mask_vars.append(mask)
@@ -273,7 +286,7 @@ def trainable_weights(self):
 
   @property
   def non_trainable_weights(self):
-    return self.layer.non_trainable_weights
+    return self.layer.non_trainable_weights + self._non_trainable_weights
 
   @property
   def updates(self):