Merge pull request #597 from johan-gras:up1

PiperOrigin-RevId: 368159207

Author: tensorflower-gardener

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

@@ -15,6 +15,8 @@
 """Tests for keras ClusterWeights wrapper API."""
 
 import itertools
+import os
+import tempfile
 
 from absl.testing import parameterized
 import tensorflow as tf
@@ -230,63 +232,46 @@ 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 weights of cluster_wrapper are the same after serialisation."""
     # Create a dummy layer for this test
     input_shape = (1, 2,)
-    l = cluster_wrapper.ClusterWeights(
+    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)
+    model = keras.Sequential([original_layer])
+
+    # 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(model, keras_file)
+      with cluster.cluster_scope():
+        loaded_layer = keras.models.load_model(keras_file).layers[0]
+
+    def assert_list_of_variables_all_equal(l1, l2):
+      self.assertLen(
+          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:
+        self.assertIn(var2.name, name_to_var_from_l1)
+        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
+    assert_list_of_variables_all_equal(original_layer.trainable_weights,
+                                       loaded_layer.trainable_weights)
+    assert_list_of_variables_all_equal(original_layer.non_trainable_weights,
+                                       loaded_layer.non_trainable_weights)
 
 if __name__ == '__main__':
   test.main()