Fix for the bug: weights/bias name should be the same for original and stripped model.

Change-Id: Ib81f4598356cd96af224ed29e9391d2f26b6bc58

Author: wwwind

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

@@ -197,6 +197,7 @@ def _strip_clustering_wrapper(layer):
       for i in range(len(layer.restore)):
         # This is why we used integers as keys
         name, weight = layer.restore[i]
+        weight_name = layer.layer.name + "/" + name
         # In both cases we use k.batch_get_value since we need physical copies
         # of the arrays to initialize a new tensor
         if i in layer.gone_variables:
@@ -209,7 +210,7 @@ def _strip_clustering_wrapper(layer):
           new_weight_value = k.batch_get_value([weight])[0]
         setattr(layer.layer,
                 name,
-                k.variable(new_weight_value, name=name))
+                k.variable(new_weight_value, name=weight_name))
       # When all weights are filled with the values, just return the underlying
       # layer since it is now fully autonomous from its wrapper
       return layer.layer

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

@@ -172,6 +172,10 @@ def testValuesAreClusteredAfterStripping(self,
     original_model = tf.keras.Sequential([
         layers.Dense(32, input_shape=(10,)),
     ])
+
+    weights_name = original_model.layers[0].weights[0].name
+    bias_name = original_model.layers[0].weights[1].name
+
     clustered_model = cluster.cluster_weights(
         original_model,
         number_of_clusters=number_of_clusters,
@@ -186,6 +190,9 @@ def testValuesAreClusteredAfterStripping(self,
     # Make sure that the stripped layer is the Dense one
     self.assertIsInstance(stripped_model.layers[0], layers.Dense)
 
+    # Check that we keep names for weights/bias
+    self.assertEqual(stripped_model.layers[0].weights[0].name, weights_name)
+    self.assertEqual(stripped_model.layers[0].weights[1].name, bias_name)
 
 if __name__ == '__main__':
   test.main()