Use TF API instead of directly importing internal members when possible

Fixed formatting to comply with python style guide

PiperOrigin-RevId: 371043407

Author: teijeong

tensorflow_model_optimization/python/core/quantization/keras/default_8bit/default_8bit_transforms.py

@@ -181,10 +181,10 @@ def _replace(self, bn_layer_node, conv_layer_node):
     if _has_custom_quantize_config(bn_layer_node, conv_layer_node):
       return bn_layer_node
 
-    conv_layer_node.layer['config']['activation'] = \
-      keras.activations.serialize(quantize_aware_activation.NoOpActivation())
-    bn_layer_node.metadata['quantize_config'] = \
-      default_8bit_quantize_configs.Default8BitOutputQuantizeConfig()
+    conv_layer_node.layer['config']['activation'] = (
+        keras.activations.serialize(quantize_aware_activation.NoOpActivation()))
+    bn_layer_node.metadata['quantize_config'] = (
+        default_8bit_quantize_configs.Default8BitOutputQuantizeConfig())
 
     return bn_layer_node
 
@@ -235,10 +235,10 @@ def _replace(self, relu_layer_node, bn_layer_node, conv_layer_node):
         relu_layer_node, bn_layer_node, conv_layer_node):
       return relu_layer_node
 
-    conv_layer_node.layer['config']['activation'] = \
-      keras.activations.serialize(quantize_aware_activation.NoOpActivation())
-    bn_layer_node.metadata['quantize_config'] = \
-      default_8bit_quantize_configs.NoOpQuantizeConfig()
+    conv_layer_node.layer['config']['activation'] = (
+        keras.activations.serialize(quantize_aware_activation.NoOpActivation()))
+    bn_layer_node.metadata['quantize_config'] = (
+        default_8bit_quantize_configs.NoOpQuantizeConfig())
 
     return relu_layer_node
 
@@ -508,8 +508,8 @@ def replacement(self, match_layer):
     relu_layer_node = match_layer
     add_layer_node = relu_layer_node.input_layers[0]
 
-    add_layer_node.metadata['quantize_config'] = \
-      default_8bit_quantize_configs.NoOpQuantizeConfig()
+    add_layer_node.metadata['quantize_config'] = (
+        default_8bit_quantize_configs.NoOpQuantizeConfig())
 
     return match_layer
 
@@ -585,8 +585,8 @@ def replacement(self, match_layer):
     concat_layer_node = match_layer
     feeding_layer_nodes = match_layer.input_layers
 
-    default_registry = default_8bit_quantize_registry.\
-        Default8BitQuantizeRegistry()
+    default_registry = (
+        default_8bit_quantize_registry.Default8BitQuantizeRegistry())
 
     feed_quantize_configs = []
     for feed_layer_node in feeding_layer_nodes:
@@ -599,8 +599,8 @@ def replacement(self, match_layer):
 
         if layer_class == keras.layers.Concatenate:
           # Input layer to Concat is also Concat. Don't quantize it.
-          feed_layer_node.metadata['quantize_config'] = \
-            default_8bit_quantize_configs.NoOpQuantizeConfig()
+          feed_layer_node.metadata['quantize_config'] = (
+              default_8bit_quantize_configs.NoOpQuantizeConfig())
           continue
 
         if not default_registry._is_supported_layer(layer_class):
@@ -619,8 +619,8 @@ def replacement(self, match_layer):
       self._disable_output_quantize(quantize_config)
 
     if not concat_layer_node.metadata.get('quantize_config'):
-      concat_layer_node.metadata['quantize_config'] = \
-        default_8bit_quantize_configs.Default8BitOutputQuantizeConfig()
+      concat_layer_node.metadata['quantize_config'] = (
+          default_8bit_quantize_configs.Default8BitOutputQuantizeConfig())
 
     return concat_layer_node
 

tensorflow_model_optimization/python/core/sparsity/keras/BUILD

@@ -131,15 +131,14 @@ py_strict_test(
     ],
 )
 
-py_library(
+py_strict_library(
     name = "pruning_impl",
     srcs = ["pruning_impl.py"],
     srcs_version = "PY3",
     visibility = ["//visibility:public"],
     deps = [
         ":pruning_utils",
         # tensorflow dep1,
-        # python:summary tensorflow dep2,
         "//tensorflow_model_optimization/python/core/keras:compat",
     ],
 )
@@ -155,16 +154,14 @@ py_strict_library(
     ],
 )
 
-py_library(
+py_strict_library(
     name = "estimator_utils",
     srcs = ["estimator_utils.py"],
     srcs_version = "PY3",
     visibility = ["//visibility:public"],
     deps = [
         ":pruning_wrapper",
-        # python:control_flow_ops tensorflow dep2,
-        # python:math_ops tensorflow dep2,
-        # python:state_ops tensorflow dep2,
+        # tensorflow dep1,
         # python/framework:for_generated_wrappers tensorflow dep2,
     ],
 )

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

@@ -14,22 +14,13 @@
 # ==============================================================================
 """Utility functions for making pruning wrapper work with estimators."""
 
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-# import g3
+import tensorflow as tf
 
-from tensorflow.python.estimator.model_fn import EstimatorSpec
-from tensorflow.python.framework import dtypes
 from tensorflow.python.framework import ops
-from tensorflow.python.ops import control_flow_ops
-from tensorflow.python.ops import math_ops
-from tensorflow.python.ops import state_ops
-from tensorflow.python.training import monitored_session
 from tensorflow_model_optimization.python.core.sparsity.keras.pruning_wrapper import PruneLowMagnitude
 
 
-class PruningEstimatorSpec(EstimatorSpec):
+class PruningEstimatorSpec(tf.estimator.EstimatorSpec):
   """Returns an EstimatorSpec modified to prune the model while training."""
 
   def __new__(cls, model, step=None, train_op=None, **kwargs):
@@ -49,13 +40,12 @@ def _get_step_increment_ops(model, step=None):
         if isinstance(layer, PruneLowMagnitude):
           if step is None:
             # Add ops to increment the pruning_step by 1
-            increment_ops.append(state_ops.assign_add(layer.pruning_step, 1))
+            increment_ops.append(tf.assign_add(layer.pruning_step, 1))
           else:
             increment_ops.append(
-                state_ops.assign(layer.pruning_step,
-                                 math_ops.cast(step, dtypes.int32)))
+                tf.assign(layer.pruning_step, tf.cast(step, tf.int32)))
 
-      return control_flow_ops.group(increment_ops)
+      return tf.group(increment_ops)
 
     pruning_ops = []
     # Grab the ops to update pruning step in every prunable layer
@@ -64,21 +54,21 @@ def _get_step_increment_ops(model, step=None):
     # Grab the model updates.
     pruning_ops.append(model.updates)
 
-    kwargs["train_op"] = control_flow_ops.group(pruning_ops, train_op)
+    kwargs["train_op"] = tf.group(pruning_ops, train_op)
 
     def init_fn(scaffold, session):  # pylint: disable=unused-argument
       return session.run(step_increment_ops)
 
     def get_new_scaffold(old_scaffold):
       if old_scaffold.init_fn is None:
-        return monitored_session.Scaffold(
+        return tf.compat.v1.train.Scaffold(
             init_fn=init_fn, copy_from_scaffold=old_scaffold)
       # TODO(suyoggupta): Figure out a way to merge the init_fn of the
       # original scaffold with the one defined above.
       raise ValueError("Scaffold provided to PruningEstimatorSpec must not "
                        "set an init_fn.")
 
-    scaffold = monitored_session.Scaffold(init_fn=init_fn)
+    scaffold = tf.compat.v1.train.Scaffold(init_fn=init_fn)
     if "scaffold" in kwargs:
       scaffold = get_new_scaffold(kwargs["scaffold"])
 

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

@@ -20,8 +20,6 @@
 
 import tensorflow as tf
 
-from tensorflow.python.ops import summary_ops_v2
-from tensorflow.python.summary import summary as summary_ops_v1
 from tensorflow_model_optimization.python.core.keras import compat as tf_compat
 from tensorflow_model_optimization.python.core.sparsity.keras import pruning_utils
 
@@ -255,9 +253,9 @@ def update_distributed():
   def add_pruning_summaries(self):
     """Adds summaries of weight sparsities and thresholds."""
     # b/(139939526): update to use public API.
-    summary = summary_ops_v1
-    if tf.executing_eagerly():
-      summary = summary_ops_v2
+    summary = tf.summary
+    if not tf.executing_eagerly():
+      summary = tf.compat.v1.summary
     summary.scalar('sparsity', self._pruning_schedule(self._step_fn())[1])
     for _, mask, threshold in self._pruning_vars:
       summary.scalar(mask.name + '/sparsity', 1.0 - tf.math.reduce_mean(mask))