Use tf.keras instead of importing keras directly

PiperOrigin-RevId: 371042274

Author: teijeong

tensorflow_model_optimization/python/core/clustering/keras/BUILD

@@ -19,7 +19,7 @@ py_strict_library(
     ],
 )
 
-py_library(
+py_strict_library(
     name = "cluster",
     srcs = ["cluster.py"],
     srcs_version = "PY3",
@@ -38,7 +38,7 @@ py_strict_library(
     visibility = ["//visibility:public"],
 )
 
-py_library(
+py_strict_library(
     name = "clustering_registry",
     srcs = ["clustering_registry.py"],
     srcs_version = "PY3",
@@ -73,7 +73,7 @@ py_strict_library(
     ],
 )
 
-py_library(
+py_strict_library(
     name = "cluster_wrapper",
     srcs = ["cluster_wrapper.py"],
     srcs_version = "PY3",
@@ -99,7 +99,7 @@ py_strict_library(
     ],
 )
 
-py_library(
+py_strict_library(
     name = "clustering_callbacks",
     srcs = ["clustering_callbacks.py"],
     srcs_version = "PY3",

tensorflow_model_optimization/python/core/clustering/keras/cluster.py

@@ -14,15 +14,15 @@
 # ==============================================================================
 """Clustering API functions for Keras models."""
 
-from tensorflow import keras
+import tensorflow as tf
 
 from tensorflow_model_optimization.python.core.clustering.keras import cluster_wrapper
 from tensorflow_model_optimization.python.core.clustering.keras import clustering_centroids
 
-k = keras.backend
-CustomObjectScope = keras.utils.CustomObjectScope
-Layer = keras.layers.Layer
-InputLayer = keras.layers.InputLayer
+k = tf.keras.backend
+CustomObjectScope = tf.keras.utils.CustomObjectScope
+Layer = tf.keras.layers.Layer
+InputLayer = tf.keras.layers.InputLayer
 
 
 def cluster_scope():
@@ -38,10 +38,10 @@ def cluster_scope():
 
   ```python
   clustered_model = cluster_weights(model, **self.params)
-  keras.models.save_model(clustered_model, keras_file)
+  tf.keras.models.save_model(clustered_model, keras_file)
 
   with cluster_scope():
-    loaded_model = keras.models.load_model(keras_file)
+    loaded_model = tf.keras.models.load_model(keras_file)
   ```
   """
   return CustomObjectScope(
@@ -92,7 +92,7 @@ def cluster_weights(to_cluster,
     'cluster_centroids_init': CentroidInitialization.DENSITY_BASED
   }
 
-  model = keras.Sequential([
+  model = tf.keras.Sequential([
       layers.Dense(10, activation='relu', input_shape=(100,)),
       cluster_weights(layers.Dense(2, activation='tanh'), **clustering_params)
   ])
@@ -174,7 +174,7 @@ def _cluster_weights(to_cluster, number_of_clusters, cluster_centroids_init,
     'preserve_sparsity': False
   }
 
-  model = keras.Sequential([
+  model = tf.keras.Sequential([
       layers.Dense(10, activation='relu', input_shape=(100,)),
       cluster_weights(layers.Dense(2, activation='tanh'), **clustering_params)
   ])
@@ -189,7 +189,7 @@ def _cluster_weights(to_cluster, number_of_clusters, cluster_centroids_init,
     'preserve_sparsity': True
   }
 
-  model = keras.Sequential([
+  model = tf.keras.Sequential([
       layers.Dense(10, activation='relu', input_shape=(100,)),
       cluster_weights(layers.Dense(2, activation='tanh'), **clustering_params)
   ])
@@ -221,16 +221,16 @@ def _cluster_weights(to_cluster, number_of_clusters, cluster_centroids_init,
         cluster_centroids_init))
 
   def _add_clustering_wrapper(layer):
-    if isinstance(layer, keras.Model):
+    if isinstance(layer, tf.keras.Model):
       # Check whether the model is a subclass.
       # NB: This check is copied from keras.py file in tensorflow.
       # There is no available public API to do this check.
       # pylint: disable=protected-access
       if (not layer._is_graph_network and
-          not isinstance(layer, keras.models.Sequential)):
+          not isinstance(layer, tf.keras.models.Sequential)):
         raise ValueError('Subclassed models are not supported currently.')
 
-      return keras.models.clone_model(
+      return tf.keras.models.clone_model(
           layer, input_tensors=None, clone_function=_add_clustering_wrapper)
     if isinstance(layer, cluster_wrapper.ClusterWeights):
       return layer
@@ -248,10 +248,9 @@ def _wrap_list(layers):
 
     return output
 
-  if isinstance(to_cluster, keras.Model):
-    return keras.models.clone_model(to_cluster,
-                                    input_tensors=None,
-                                    clone_function=_add_clustering_wrapper)
+  if isinstance(to_cluster, tf.keras.Model):
+    return tf.keras.models.clone_model(
+        to_cluster, input_tensors=None, clone_function=_add_clustering_wrapper)
   if isinstance(to_cluster, Layer):
     return _add_clustering_wrapper(layer=to_cluster)
   if isinstance(to_cluster, list):
@@ -285,13 +284,13 @@ def strip_clustering(model):
   ```
   The exported_model and the orig_model have the same structure.
   """
-  if not isinstance(model, keras.Model):
+  if not isinstance(model, tf.keras.Model):
     raise ValueError(
         'Expected model to be a `tf.keras.Model` instance but got: ', model)
 
   def _strip_clustering_wrapper(layer):
-    if isinstance(layer, keras.Model):
-      return keras.models.clone_model(
+    if isinstance(layer, tf.keras.Model):
+      return tf.keras.models.clone_model(
           layer, input_tensors=None, clone_function=_strip_clustering_wrapper)
 
     elif isinstance(layer, cluster_wrapper.ClusterWeights):
@@ -317,6 +316,5 @@ def _strip_clustering_wrapper(layer):
     return layer
 
   # Just copy the model with the right callback
-  return keras.models.clone_model(model,
-                                  input_tensors=None,
-                                  clone_function=_strip_clustering_wrapper)
+  return tf.keras.models.clone_model(
+      model, input_tensors=None, clone_function=_strip_clustering_wrapper)

tensorflow_model_optimization/python/core/clustering/keras/cluster_wrapper.py

@@ -16,8 +16,6 @@
 
 import tensorflow as tf
 
-from tensorflow.keras import initializers
-
 from tensorflow_model_optimization.python.core.clustering.keras import cluster_config
 from tensorflow_model_optimization.python.core.clustering.keras import clusterable_layer
 from tensorflow_model_optimization.python.core.clustering.keras import clustering_centroids
@@ -130,8 +128,8 @@ def __init__(self,
     # If the input shape was specified, then we need to preserve this
     # information in the layer. If this info is not preserved, then the `built`
     # state will not be preserved between serializations.
-    if (not hasattr(self, '_batch_input_shape')
-        and hasattr(layer, '_batch_input_shape')):
+    if (not hasattr(self, '_batch_input_shape') and
+        hasattr(layer, '_batch_input_shape')):
       self._batch_input_shape = self.layer._batch_input_shape
 
     # Save the input shape specified in the build
@@ -171,7 +169,7 @@ def build(self, input_shape):
           shape=(self.number_of_clusters,),
           dtype=weight.dtype,
           trainable=True,
-          initializer=initializers.Constant(value=cluster_centroids))
+          initializer=tf.keras.initializers.Constant(value=cluster_centroids))
 
       # Init the weight clustering algorithm
       self.clustering_algorithms[weight_name] = (
@@ -192,7 +190,7 @@ def build(self, input_shape):
           trainable=False,
           synchronization=tf.VariableSynchronization.ON_READ,
           aggregation=tf.VariableAggregation.ONLY_FIRST_REPLICA,
-          initializer=initializers.Constant(value=pulling_indices))
+          initializer=tf.keras.initializers.Constant(value=pulling_indices))
 
       if self.preserve_sparsity:
         # Init the sparsity mask

tensorflow_model_optimization/python/core/clustering/keras/clustering_callbacks.py

@@ -16,11 +16,10 @@
 
 import tensorflow as tf
 
-from tensorflow import keras
 from tensorflow_model_optimization.python.core.keras import compat
 
 
-class ClusteringSummaries(keras.callbacks.TensorBoard):
+class ClusteringSummaries(tf.keras.callbacks.TensorBoard):
   """Helper class to create tensorboard summaries for the clustering progress.
 
     This class is derived from tf.keras.callbacks.TensorBoard and just adds

tensorflow_model_optimization/python/core/clustering/keras/clustering_registry.py

@@ -15,11 +15,11 @@
 """Registry responsible for built-in keras classes."""
 
 import tensorflow as tf
-from tensorflow.keras import layers
 
 from tensorflow_model_optimization.python.core.clustering.keras import clusterable_layer
 from tensorflow_model_optimization.python.core.clustering.keras import clustering_algorithm
 
+layers = tf.keras.layers
 AbstractClusteringAlgorithm = clustering_algorithm.AbstractClusteringAlgorithm
 
 

tensorflow_model_optimization/python/core/quantization/keras/collaborative_optimizations/cluster_preserve/BUILD

@@ -28,14 +28,13 @@ py_strict_library(
     ],
 )
 
-py_library(
+py_strict_library(
     name = "cluster_preserve_quantize_registry",
     srcs = [
         "cluster_preserve_quantize_registry.py",
     ],
     srcs_version = "PY3",
     deps = [
-        ":cluster_utils",
         # tensorflow dep1,
         "//tensorflow_model_optimization/python/core/clustering/keras:clustering_registry",
         "//tensorflow_model_optimization/python/core/quantization/keras:quant_ops",

tensorflow_model_optimization/python/core/quantization/keras/collaborative_optimizations/cluster_preserve/cluster_preserve_quantize_registry.py

@@ -15,14 +15,14 @@
 """Registry responsible for built-in keras classes."""
 
 import tensorflow as tf
-from tensorflow.keras import backend as K
 
 from tensorflow_model_optimization.python.core.clustering.keras import clustering_registry
 from tensorflow_model_optimization.python.core.quantization.keras import quant_ops
 from tensorflow_model_optimization.python.core.quantization.keras import quantizers
 from tensorflow_model_optimization.python.core.quantization.keras.default_8bit import default_8bit_quantize_registry
 from tensorflow_model_optimization.python.core.quantization.keras.default_8bit import default_8bit_quantizers
 
+K = tf.keras.backend
 layers = tf.keras.layers