Add support for tf.distribute after enabling updates of cluster indices

Author: Ruomei

tensorflow_model_optimization/python/core/clustering/keras/BUILD

@@ -150,3 +150,15 @@ py_test(
         "//tensorflow_model_optimization/python/core/keras:compat",
     ],
 )
+
+py_test(
+    name = "cluster_distributed_test",
+    srcs = ["cluster_distributed_test.py"],
+    python_version = "PY3",
+    visibility = ["//visibility:public"],
+    deps = [
+        ":cluster",
+        "//tensorflow_model_optimization/python/core/keras:test_utils",
+        # tensorflow dep1,
+    ],
+)

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

@@ -48,7 +48,7 @@ def cluster_scope():
   """
   return CustomObjectScope(
       {
-          'ClusterWeights' : cluster_wrapper.ClusterWeights
+          'ClusterWeights': cluster_wrapper.ClusterWeights
       }
   )
 

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

@@ -0,0 +1,138 @@
+# Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Distributed clustering test."""
+
+from absl.testing import parameterized
+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
+
+keras = tf.keras
+CentroidInitialization = cluster_config.CentroidInitialization
+
+
+def _distribution_strategies():
+  return [
+      tf.distribute.MirroredStrategy()
+  ]
+
+
+class ClusterDistributedTest(tf.test.TestCase, parameterized.TestCase):
+  """Distributed tests for clustering."""
+
+  def setUp(self):
+    super(ClusterDistributedTest, self).setUp()
+    self.params = {
+        "number_of_clusters": 2,
+        "cluster_centroids_init": CentroidInitialization.LINEAR
+    }
+
+
+  @parameterized.parameters(_distribution_strategies())
+  def testClusterSimpleDenseModel(self, distribution):
+    """End-to-end test."""
+    with distribution.scope():
+      model = cluster.cluster_weights(
+          keras_test_utils.build_simple_dense_model(), **self.params)
+      model.compile(
+          loss='categorical_crossentropy',
+          optimizer='sgd',
+          metrics=['accuracy'])
+
+    model.summary()
+    model.fit(
+        np.random.rand(20, 10),
+        keras.utils.to_categorical(np.random.randint(5, size=(20, 1)), 5),
+        epochs=1,
+        batch_size=20)
+    model.predict(np.random.rand(20, 10))
+
+    stripped_model = cluster.strip_clustering(model)
+    weights_as_list = stripped_model.get_weights()[0].reshape(-1,).tolist()
+    unique_weights = set(weights_as_list)
+    self.assertLessEqual(len(unique_weights), self.params["number_of_clusters"])
+
+  @parameterized.parameters(_distribution_strategies())
+  def testAssociationValuesPerReplica(self, distribution):
+    """Verifies that associations of weights are updated per replica."""
+    assert tf.distribute.get_replica_context() is not None
+    with distribution.scope():
+      assert tf.distribute.get_replica_context() is None
+      input_shape = (1, 2)
+      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"]
+      )
+      l.build(input_shape)
+
+      clusterable_weights = l.layer.get_clusterable_weights()
+      self.assertEqual(len(clusterable_weights), 1)
+      weights_name = clusterable_weights[0][0]
+      self.assertEqual(weights_name, 'kernel')
+      centroids1 = l.cluster_centroids_tf[weights_name]
+
+      mean_weight = tf.reduce_mean(l.layer.kernel)
+      min_weight = tf.reduce_min(l.layer.kernel)
+      max_weight = tf.reduce_max(l.layer.kernel)
+      max_dist = max_weight - min_weight
+
+      def assert_all_cluster_indices(per_replica, indices_val):
+        if indices_val == 1:
+          val_tensor = tf.dtypes.cast(
+              tf.ones(shape=output_shape), per_replica[0].dtype)
+        if indices_val == 0:
+          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
+              )
+          )
+          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)
+
+      centroids1 = distribution.extended.update(
+          centroids1, update_fn, args=(initial_val,))
+      l.call(tf.ones(shape=input_shape))
+
+      clst_indices = l.pulling_indices_tf[weights_name]
+      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)
+      centroids2 = l.cluster_centroids_tf[weights_name]
+      centroids2 = distribution.extended.update(
+          centroids2, update_fn, args=(second_val,))
+      l.call(tf.ones(shape=input_shape))
+
+      clst_indices = l.pulling_indices_tf[weights_name]
+      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_wrapper.py

@@ -222,6 +222,8 @@ def build(self, input_shape):
           shape=pulling_indices.shape,
           dtype=tf.int32,
           trainable=False,
+          synchronization=tf.VariableSynchronization.ON_READ,
+          aggregation=tf.VariableAggregation.ONLY_FIRST_REPLICA,
           initializer=initializers.Constant(
               value=k.batch_get_value([pulling_indices])[0]
           )
@@ -254,7 +256,7 @@ def fn():
     # This loop stores pairs of weight names and how to restore them
     for ct, weight in enumerate(self.layer.weights):
       name = self._weight_name(weight.name)
-      full_name = self.layer.name + "/" + name
+      full_name = '{}/{}'.format(self.layer.name, name)
       if ct in self.gone_variables:
         # Again, not sure if this is needed
         weight_name = clusterable_weights_to_variables[name]

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

@@ -247,5 +247,67 @@ def assert_all_weights_associated(weights, centroid_index):
     assert_all_weights_associated(l.layer.kernel, centroid_index=1)
 
 
+  def testClusterReassociation(self):
+    """
+    Verifies that the association of weights to cluster centroids are updated
+    every iteration.
+    """
+
+    # Create a dummy layer for this test
+    input_shape = (1, 2,)
+    l = cluster_wrapper.ClusterWeights(
+        keras.layers.Dense(8, input_shape=input_shape),
+        number_of_clusters=2,
+        cluster_centroids_init=CentroidInitialization.LINEAR
+    )
+    # Build a layer with the given shape
+    l.build(input_shape)
+
+    # Get name of the clusterable weights
+    clusterable_weights = l.layer.get_clusterable_weights()
+    self.assertEqual(len(clusterable_weights), 1)
+    weights_name = clusterable_weights[0][0]
+    self.assertEqual(weights_name, 'kernel')
+    # Get cluster centroids
+    centroids = l.cluster_centroids_tf[weights_name]
+
+    # Calculate some statistics of the weights to set the centroids later on
+    mean_weight = tf.reduce_mean(l.layer.kernel)
+    min_weight = tf.reduce_min(l.layer.kernel)
+    max_weight = tf.reduce_max(l.layer.kernel)
+    max_dist = max_weight - min_weight
+
+    def assert_all_weights_associated(weights, centroid_index):
+      """Helper function to make sure that all weights are associated with one
+      centroid."""
+      all_associated = tf.reduce_all(
+          tf.equal(
+              weights,
+              tf.constant(centroids[centroid_index], shape=weights.shape)
+          )
+      )
+      self.assertTrue(all_associated)
+
+    # Set centroids so that all weights should be re-associated with centroid 0
+    centroids[0].assign(mean_weight)
+    centroids[1].assign(mean_weight + 2.0 * max_dist)
+
+    # Update associations of weights to centroids
+    l.call(tf.ones(shape=input_shape))
+
+    # Weights should now be all clustered with the centroid 0
+    assert_all_weights_associated(l.layer.kernel, centroid_index=0)
+
+    # Set centroids so that all weights should be re-associated with centroid 1
+    centroids[0].assign(mean_weight - 2.0 * max_dist)
+    centroids[1].assign(mean_weight)
+
+    # Update associations of weights to centroids
+    l.call(tf.ones(shape=input_shape))
+
+    # Weights should now be all clustered with the centroid 1
+    assert_all_weights_associated(l.layer.kernel, centroid_index=1)
+
+
 if __name__ == '__main__':
   test.main()