Fix formatting for clustering codes

* Ran automated formatting fix
* Eliminated pylint errors

PiperOrigin-RevId: 370812580

Author: teijeong

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

@@ -299,8 +299,10 @@ def _strip_clustering_wrapper(layer):
       layer.update_clustered_weights_associations()
 
       # Construct a list of weights to initialize the clean layer
-      updated_weights = layer.layer.get_weights()  # non clusterable weights only
-      for position_variable, weight_name in layer.position_original_weights.items():
+      # non-clusterable weights only
+      updated_weights = layer.layer.get_weights()
+      for (position_variable,
+           weight_name) in layer.position_original_weights.items():
         # Add the clustered weights at the correct position
         clustered_weight = getattr(layer.layer, weight_name)
         updated_weights.insert(position_variable, clustered_weight)

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

@@ -18,19 +18,17 @@
 import numpy as np
 import tensorflow as tf
 
-from tensorflow_model_optimization.python.core.keras import test_utils as keras_test_utils
 from tensorflow_model_optimization.python.core.clustering.keras import cluster
 from tensorflow_model_optimization.python.core.clustering.keras import cluster_config
 from tensorflow_model_optimization.python.core.clustering.keras import cluster_wrapper
+from tensorflow_model_optimization.python.core.keras import test_utils as keras_test_utils
 
 keras = tf.keras
 CentroidInitialization = cluster_config.CentroidInitialization
 
 
 def _distribution_strategies():
-  return [
-      tf.distribute.MirroredStrategy()
-  ]
+  return [tf.distribute.MirroredStrategy()]
 
 
 class ClusterDistributedTest(tf.test.TestCase, parameterized.TestCase):
@@ -39,8 +37,8 @@ class ClusterDistributedTest(tf.test.TestCase, parameterized.TestCase):
   def setUp(self):
     super(ClusterDistributedTest, self).setUp()
     self.params = {
-        "number_of_clusters": 2,
-        "cluster_centroids_init": CentroidInitialization.LINEAR
+        'number_of_clusters': 2,
+        'cluster_centroids_init': CentroidInitialization.LINEAR
     }
 
   @parameterized.parameters(_distribution_strategies())
@@ -63,9 +61,10 @@ def testClusterSimpleDenseModel(self, distribution):
     model.predict(np.random.rand(20, 10))
 
     stripped_model = cluster.strip_clustering(model)
-    weights_as_list = stripped_model.layers[0].kernel.numpy().reshape(-1,).tolist()
+    weights_as_list = stripped_model.layers[0].kernel.numpy().reshape(
+        -1,).tolist()
     unique_weights = set(weights_as_list)
-    self.assertLessEqual(len(unique_weights), self.params["number_of_clusters"])
+    self.assertLessEqual(len(unique_weights), self.params['number_of_clusters'])
 
   @parameterized.parameters(_distribution_strategies())
   def testAssociationValuesPerReplica(self, distribution):
@@ -77,13 +76,12 @@ def testAssociationValuesPerReplica(self, distribution):
       output_shape = (2, 8)
       l = cluster_wrapper.ClusterWeights(
           keras.layers.Dense(8, input_shape=input_shape),
-          number_of_clusters=self.params["number_of_clusters"],
-          cluster_centroids_init=self.params["cluster_centroids_init"]
-      )
+          number_of_clusters=self.params['number_of_clusters'],
+          cluster_centroids_init=self.params['cluster_centroids_init'])
       l.build(input_shape)
 
       clusterable_weights = l.layer.get_clusterable_weights()
-      self.assertEqual(len(clusterable_weights), 1)
+      self.assertLen(clusterable_weights, 1)
       weights_name = clusterable_weights[0][0]
       self.assertEqual(weights_name, 'kernel')
       centroids1 = l.cluster_centroids[weights_name]
@@ -101,18 +99,14 @@ def assert_all_cluster_indices(per_replica, indices_val):
           val_tensor = tf.dtypes.cast(
               tf.zeros(shape=output_shape), per_replica[0].dtype)
         for i in range(0, len(per_replica)):
-          all_equal = tf.reduce_all(
-              tf.equal(
-                  per_replica[i], val_tensor
-              )
-          )
+          all_equal = tf.reduce_all(tf.equal(per_replica[i], val_tensor))
           self.assertTrue(all_equal)
 
       def update_fn(v, val):
         return v.assign(val)
 
-      initial_val = tf.Variable([mean_weight, mean_weight + 2.0 * max_dist], \
-        aggregation=tf.VariableAggregation.MEAN)
+      initial_val = tf.Variable([mean_weight, mean_weight + 2.0 * max_dist],
+                                aggregation=tf.VariableAggregation.MEAN)
 
       centroids1 = distribution.extended.update(
           centroids1, update_fn, args=(initial_val,))
@@ -122,8 +116,8 @@ def update_fn(v, val):
       per_replica = distribution.experimental_local_results(clst_indices)
       assert_all_cluster_indices(per_replica, 0)
 
-      second_val = tf.Variable([mean_weight - 2.0 * max_dist, mean_weight], \
-        aggregation=tf.VariableAggregation.MEAN)
+      second_val = tf.Variable([mean_weight - 2.0 * max_dist, mean_weight],
+                               aggregation=tf.VariableAggregation.MEAN)
       centroids2 = l.cluster_centroids[weights_name]
       centroids2 = distribution.extended.update(
           centroids2, update_fn, args=(second_val,))
@@ -133,5 +127,6 @@ def update_fn(v, val):
       per_replica = distribution.experimental_local_results(clst_indices)
       assert_all_cluster_indices(per_replica, 1)
 
+
 if __name__ == '__main__':
   tf.test.main()

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

@@ -152,7 +152,8 @@ def testValuesRemainClusteredAfterTraining(self):
 
   @keras_parameterized.run_all_keras_modes
   def testSparsityIsPreservedDuringTraining(self):
-    """Set a specific random seed to ensure that we get some null weights to test sparsity preservation with."""
+    # Set a specific random seed to ensure that we get some null weights to
+    # test sparsity preservation with.
     tf.random.set_seed(1)
 
     # Verifies that training a clustered model does not destroy the sparsity of
@@ -187,7 +188,8 @@ def testSparsityIsPreservedDuringTraining(self):
     stripped_model_after_tuning = cluster.strip_clustering(clustered_model)
     weights_after_tuning = stripped_model_after_tuning.layers[0].kernel
     non_zero_weight_indices_after_tuning = np.nonzero(weights_after_tuning)
-    weights_as_list_after_tuning = weights_after_tuning.numpy().reshape(-1,).tolist()
+    weights_as_list_after_tuning = weights_after_tuning.numpy().reshape(
+        -1,).tolist()
     unique_weights_after_tuning = set(weights_as_list_after_tuning)
 
     # Check that the null weights stayed the same before and after tuning.
@@ -245,15 +247,15 @@ def testEndToEndDeepLayer(self):
 
     def clusters_check(stripped_model):
       # inner dense layer
-      weights_as_list = stripped_model.submodules[1].trainable_weights[0].\
-        numpy().flatten()
+      weights_as_list = (
+          stripped_model.submodules[1].trainable_weights[0].numpy().flatten())
       unique_weights = set(weights_as_list)
       self.assertLessEqual(
           len(unique_weights), self.params["number_of_clusters"])
 
       # outer dense layer
-      weights_as_list = stripped_model.submodules[4].trainable_weights[0].\
-        numpy().flatten()
+      weights_as_list = (
+          stripped_model.submodules[4].trainable_weights[0].numpy().flatten())
       unique_weights = set(weights_as_list)
       self.assertLessEqual(
           len(unique_weights), self.params["number_of_clusters"])
@@ -276,23 +278,22 @@ def testEndToEndDeepLayer2(self):
 
     def clusters_check(stripped_model):
       # first inner dense layer
-      weights_as_list = stripped_model.submodules[1].trainable_weights[0].\
-        numpy().flatten()
+      weights_as_list = (
+          stripped_model.submodules[1].trainable_weights[0].numpy().flatten())
       unique_weights = set(weights_as_list)
       self.assertLessEqual(
           len(unique_weights), self.params["number_of_clusters"])
 
       # second inner dense layer
-      weights_as_list = stripped_model.submodules[4].\
-        trainable_weights[0].\
-        numpy().flatten()
+      weights_as_list = (
+          stripped_model.submodules[4].trainable_weights[0].numpy().flatten())
       unique_weights = set(weights_as_list)
       self.assertLessEqual(
           len(unique_weights), self.params["number_of_clusters"])
 
       # outer dense layer
-      weights_as_list = stripped_model.submodules[7].trainable_weights[0].\
-        numpy().flatten()
+      weights_as_list = (
+          stripped_model.submodules[7].trainable_weights[0].numpy().flatten())
       unique_weights = set(weights_as_list)
       self.assertLessEqual(
           len(unique_weights), self.params["number_of_clusters"])
@@ -301,51 +302,46 @@ def clusters_check(stripped_model):
 
   @keras_parameterized.run_all_keras_modes
   def testWeightsAreLearningDuringClustering(self):
-    """Verifies that training a clustered model does update
-    original_weights, clustered_centroids and bias."""
-    original_model = keras.Sequential([
-      layers.Dense(5, input_shape=(5,))
-    ])
+    """Verifies that weights are updated during training a clustered model.
+
+    Training a clustered model should update original_weights,
+    clustered_centroids and bias.
+    """
+    original_model = keras.Sequential([layers.Dense(5, input_shape=(5,))])
 
     clustered_model = cluster.cluster_weights(original_model, **self.params)
 
     clustered_model.compile(
-      loss=keras.losses.categorical_crossentropy,
-      optimizer="adam",
-      metrics=["accuracy"],
+        loss=keras.losses.categorical_crossentropy,
+        optimizer="adam",
+        metrics=["accuracy"],
     )
 
     class CheckWeightsCallback(keras.callbacks.Callback):
+
       def on_train_batch_begin(self, batch, logs=None):
         # Save weights before batch
         self.original_weight_kernel = (
-          self.model.layers[0].original_clusterable_weights['kernel'].numpy()
-        )
+            self.model.layers[0].original_clusterable_weights["kernel"].numpy())
         self.cluster_centroids_kernel = (
-          self.model.layers[0].cluster_centroids['kernel'].numpy()
-        )
-        self.bias = (
-          self.model.layers[0].layer.bias.numpy()
-        )
+            self.model.layers[0].cluster_centroids["kernel"].numpy())
+        self.bias = (self.model.layers[0].layer.bias.numpy())
 
       def on_train_batch_end(self, batch, logs=None):
         # Check weights are different after batch
         assert not np.array_equal(
-          self.original_weight_kernel,
-          self.model.layers[0].original_clusterable_weights['kernel'].numpy()
-        )
+            self.original_weight_kernel,
+            self.model.layers[0].original_clusterable_weights["kernel"].numpy())
         assert not np.array_equal(
-          self.cluster_centroids_kernel,
-          self.model.layers[0].cluster_centroids['kernel'].numpy()
-        )
-        assert not np.array_equal(
-          self.bias,
-          self.model.layers[0].layer.bias.numpy()
-        )
-
-    clustered_model.fit(x=self.dataset_generator(),
-                        steps_per_epoch=5,
-                        callbacks=[CheckWeightsCallback()])
+            self.cluster_centroids_kernel,
+            self.model.layers[0].cluster_centroids["kernel"].numpy())
+        assert not np.array_equal(self.bias,
+                                  self.model.layers[0].layer.bias.numpy())
+
+    clustered_model.fit(
+        x=self.dataset_generator(),
+        steps_per_epoch=5,
+        callbacks=[CheckWeightsCallback()])
 
 
 if __name__ == "__main__":