Merge pull request #517 from wwwind:bug_inconsistency_weights_name

PiperOrigin-RevId: 328666954

Author: tensorflower-gardener

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

@@ -196,7 +196,7 @@ def _strip_clustering_wrapper(layer):
       layer.layer._non_trainable_weights = []
       for i in range(len(layer.restore)):
         # This is why we used integers as keys
-        name, weight = layer.restore[i]
+        name, weight_name, weight = layer.restore[i]
         # 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 +209,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.py

@@ -244,12 +244,13 @@ def fn():
 
     for ct, weight in enumerate(self.layer.weights):
       name = self._weight_name(weight.name)
+      full_name = self.layer.name + "/" + name
       if ct in self.gone_variables:
         # Again, not sure if this is needed
         weight_name = clusterable_weights_to_variables[name]
-        self.restore.append((name, get_updater(weight_name)))
+        self.restore.append((name, full_name, get_updater(weight_name)))
       else:
-        self.restore.append((name, weight))
+        self.restore.append((name, full_name, weight))
 
   def call(self, inputs):
     # Go through all tensors and replace them with their clustered copies.

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

@@ -170,8 +170,12 @@ def testValuesAreClusteredAfterStripping(self,
     or equal to number_of_clusters.
     """
     original_model = tf.keras.Sequential([
-        layers.Dense(32, input_shape=(10,)),
+        layers.Dense(32, input_shape=(10,), name='abc'),
     ])
+
+    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()