Update the way of incrementing pruning_step.

PiperOrigin-RevId: 352878617

Author: tensorflower-gardener

tensorflow_model_optimization/python/core/sparsity/keras/BUILD

@@ -91,6 +91,7 @@ py_library(
         # numpy dep1,
         # tensorflow dep1,
         # python/keras/utils:generic_utils tensorflow dep2,
+        "//tensorflow_model_optimization/python/core/keras:compat",
         "//tensorflow_model_optimization/python/core/keras:utils",
     ],
 )

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

@@ -473,7 +473,7 @@ def testPruneStopAndRestart_PreservesSparsity(self, save_restore_fn):
     if save_restore_fn.__name__ == '_save_restore_tf_model':
       return
 
-    begin_step, end_step = 0, 4
+    begin_step, end_step = 1, 4
     params = self.params
     params['pruning_schedule'] = pruning_schedule.PolynomialDecay(
         0.2, 0.6, begin_step, end_step, 3, 1)
@@ -542,7 +542,9 @@ class PruneIntegrationCustomTrainingLoopTest(tf.test.TestCase,
 
   def testPrunesModel_CustomTrainingLoop_ReachesTargetSparsity(self):
     pruned_model = prune.prune_low_magnitude(
-        keras_test_utils.build_simple_dense_model())
+        keras_test_utils.build_simple_dense_model(),
+        pruning_schedule=pruning_schedule.ConstantSparsity(
+            target_sparsity=0.5, begin_step=0, frequency=1))
 
     batch_size = 20
     x_train = np.random.rand(20, 10)

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

@@ -64,15 +64,13 @@ def __init__(self):
   def on_train_begin(self, logs=None):
     # Collect all the prunable layers in the model.
     self.prunable_layers = _collect_prunable_layers(self.model)
-    self.step = K.get_value(self.model.optimizer.iterations)
-
-  def on_train_batch_begin(self, batch, logs=None):
-    tuples = []
-    for layer in self.prunable_layers:
-      tuples.append((layer.pruning_step, self.step))
-
-    K.batch_set_value(tuples)
-    self.step = self.step + 1
+    # If the model is newly created/initialized, set the 'pruning_step' to 0.
+    # If the model is saved and then restored, do nothing.
+    if self.prunable_layers[0].pruning_step == -1:
+      tuples = []
+      for layer in self.prunable_layers:
+        tuples.append((layer.pruning_step, 0))
+      K.batch_set_value(tuples)
 
   def on_epoch_end(self, batch, logs=None):
     # At the end of every epoch, remask the weights. This ensures that when

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

@@ -70,9 +70,9 @@ def testUpdatePruningStepsAndLogsSummaries(self):
         ])
 
     self.assertEqual(
-        2, tf.keras.backend.get_value(pruned_model.layers[0].pruning_step))
+        3, tf.keras.backend.get_value(pruned_model.layers[0].pruning_step))
     self.assertEqual(
-        2, tf.keras.backend.get_value(pruned_model.layers[1].pruning_step))
+        3, tf.keras.backend.get_value(pruned_model.layers[1].pruning_step))
 
     self._assertLogsExist(log_dir)
 
@@ -111,11 +111,23 @@ def testUpdatePruningStepsAndLogsSummaries_CustomTrainingLoop(self):
       step_callback.on_epoch_end(batch=unused_arg)
 
     self.assertEqual(
-        2, tf.keras.backend.get_value(pruned_model.layers[0].pruning_step))
+        3, tf.keras.backend.get_value(pruned_model.layers[0].pruning_step))
     self.assertEqual(
-        2, tf.keras.backend.get_value(pruned_model.layers[1].pruning_step))
+        3, tf.keras.backend.get_value(pruned_model.layers[1].pruning_step))
     self._assertLogsExist(log_dir)
 
+  @keras_parameterized.run_all_keras_modes
+  def testUpdatePruningStepsAndLogsSummaries_RunInference(self):
+    pruned_model, _, _, x_train, _ = self._pruned_model_setup(
+        custom_training_loop=True)
+    model_output = pruned_model(x_train)
+    del model_output
+
+    self.assertEqual(
+        -1, tf.keras.backend.get_value(pruned_model.layers[0].pruning_step))
+    self.assertEqual(
+        -1, tf.keras.backend.get_value(pruned_model.layers[1].pruning_step))
+
   @keras_parameterized.run_all_keras_modes
   def testPruneTrainingRaisesError_PruningStepCallbackMissing(self):
     pruned_model, x_train, y_train = self._pruned_model_setup()

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

@@ -26,9 +26,8 @@
 
 # b/(139939526): update to use public API.
 from tensorflow.python.keras.utils import generic_utils
-
+from tensorflow_model_optimization.python.core.keras import compat as tf_compat
 from tensorflow_model_optimization.python.core.keras import utils
-
 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
@@ -234,11 +233,14 @@ def call(self, inputs, training=None):
     if training is None:
       training = K.learning_phase()
 
+    def increment_step():
+      return tf_compat.assign(self.pruning_step, self.pruning_step + 1)
+
     def add_update():
       with tf.control_dependencies([
           tf.debugging.assert_greater_equal(
               self.pruning_step,
-              np.int64(0),
+              np.int64(1),
               message=self._PRUNE_CALLBACK_ERROR_MSG)
       ]):
         with tf.control_dependencies(
@@ -248,8 +250,14 @@ def add_update():
     def no_op():
       return tf.no_op('no_update')
 
-    update_op = utils.smart_cond(training, add_update, no_op)
-    self.add_update(update_op)
+    # Increment the 'pruning_step' after each step.
+    update_pruning_step = utils.smart_cond(training, increment_step, no_op)
+    self.add_update(update_pruning_step)
+
+    # Update mask tensor after each 'pruning_frequency' steps.
+    update_mask = utils.smart_cond(training, add_update, no_op)
+    self.add_update(update_mask)
+
     # Always execute the op that performs weights = weights * mask
     # Relies on UpdatePruningStep callback to ensure the weights
     # are sparse after the final backpropagation.