Update pruning api examples to use PolynomialDecay instead of ConstantSparsity.

PiperOrigin-RevId: 402817279

Author: rino20

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

@@ -14,7 +14,7 @@
 # ==============================================================================
 """Utility functions for making pruning wrapper work with estimators."""
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 from tensorflow.python.framework import ops
 from tensorflow_model_optimization.python.core.sparsity.keras.pruning_wrapper import PruneLowMagnitude
@@ -32,6 +32,12 @@ def __new__(cls, model, step=None, train_op=None, **kwargs):
       raise ValueError(
           "Must provide train_op for creating a PruningEstimatorSpec")
 
+    for layer in model.layers:
+      # If the model is newly created/initialized, set the 'pruning_step' to 0.
+      # Otherwise, do nothing.
+      if isinstance(layer, PruneLowMagnitude) and layer.pruning_step == -1:
+        tf.assign(layer.pruning_step, 0)
+
     def _get_step_increment_ops(model, step=None):
       """Returns ops to increment the pruning_step in the prunable layers."""
       increment_ops = []
@@ -43,7 +49,7 @@ def _get_step_increment_ops(model, step=None):
             increment_ops.append(tf.assign_add(layer.pruning_step, 1))
           else:
             increment_ops.append(
-                tf.assign(layer.pruning_step, tf.cast(step, tf.int32)))
+                tf.assign(layer.pruning_step, tf.cast(step, tf.int64)))
 
       return tf.group(increment_ops)
 

tensorflow_model_optimization/python/examples/sparsity/keras/mnist/mnist_cnn.py

@@ -25,7 +25,7 @@
 from tensorflow_model_optimization.python.core.sparsity.keras import pruning_callbacks
 from tensorflow_model_optimization.python.core.sparsity.keras import pruning_schedule
 
-ConstantSparsity = pruning_schedule.ConstantSparsity
+PolynomialDecay = pruning_schedule.PolynomialDecay
 keras = tf.keras
 l = keras.layers
 
@@ -159,7 +159,13 @@ def main(unused_argv):
   y_test = tf.keras.utils.to_categorical(y_test, num_classes)
 
   pruning_params = {
-      'pruning_schedule': ConstantSparsity(0.75, begin_step=2000, frequency=100)
+      'pruning_schedule':
+          PolynomialDecay(
+              initial_sparsity=0.1,
+              final_sparsity=0.75,
+              begin_step=1000,
+              end_step=5000,
+              frequency=100)
   }
 
   layerwise_model = build_layerwise_model(input_shape, **pruning_params)