Adding a unit test to check all needed weights are returned by cluster_wrapper.

Change-Id: I60b0c901b4d7b699cc44848481518df09d36a97a

Author: johan-gras

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

@@ -15,6 +15,9 @@
 """Tests for keras ClusterWeights wrapper API."""
 
 import itertools
+import tempfile
+import os
+import tensorflow as tf
 
 from absl.testing import parameterized
 import tensorflow as tf
@@ -230,63 +233,44 @@ def assert_all_weights_associated(weights, centroid_index):
     # Weights should now be all clustered with the centroid 1
     assert_all_weights_associated(l.layer.kernel, centroid_index=1)
 
-  def testClusterReassociation2(self):
-    """Verifies that the association of weights to cluster centroids are updated every iteration."""
-
+  def testSameWeightsAreReturnedBeforeAndAfterSerialisation(self):
+    """Verify that the weights of a cluster_wrapper are the same
+    before and after serialisation."""
     # 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
+    original_layer = 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.assertLen(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)
-
+    original_layer.build(input_shape)
+
+    # Save and load the layer in a temp directory
+    with tempfile.TemporaryDirectory() as tmp_dir_name:
+      keras_file = os.path.join(tmp_dir_name, 'keras_model')
+      keras.models.save_model(original_layer, keras_file)
+      with cluster.cluster_scope():
+        loaded_layer = keras.models.load_model(keras_file)
+
+    def assertListOfVariablesAllEqual(l1, l2):
+      assert len(l1) == len(l2), \
+        "lists l1 and l2 are not equal: \n l1={l1} \n l2={l2}".format(
+          l1=[v.name for v in l1],
+          l2=[v.name for v in l2])
+
+      name_to_var_from_l1 = {var.name: var for var in l1}
+      for var2 in l2:
+        arr1 = name_to_var_from_l1[var2.name].numpy()
+        arr2 = var2.numpy()
+        self.assertAllEqual(arr1, arr2)
+
+    # Check that trainable_weights and non_trainable_weights are the same
+    # in the original layer and loaded layer
+    assertListOfVariablesAllEqual(original_layer.trainable_weights,
+                                  loaded_layer.trainable_weights)
+    assertListOfVariablesAllEqual(original_layer.non_trainable_weights,
+                                  loaded_layer.non_trainable_weights)
 
 if __name__ == '__main__':
   test.main()