Add a fallback that uses (name, weight) tuple list when weight map key has a collision.

PiperOrigin-RevId: 367911645

Author: Xhark

tensorflow_model_optimization/python/core/quantization/keras/graph_transformations/model_transformer.py

@@ -95,6 +95,9 @@ def _get_layers(self, layer_names):
   def _get_layer_weights(self, layer_name):
     return self._layer_weights_map.get(layer_name, {})
 
+  def _get_layer_names_and_weights(self, layer_name):
+    return self._layer_names_and_weights_map.get(layer_name, {})
+
   def _get_layer_metadata(self, layer_name):
     return self._layer_metadata_map.get(layer_name, {})
 
@@ -191,7 +194,9 @@ def _match_layer_with_inputs(self, layer, pattern, is_head_node):
     if len(pattern.inputs) == 0:
       # Leaf layer in pattern.
       return LayerNode(layer, self._get_layer_weights(layer['config']['name']),
-                       [], self._get_layer_metadata(layer['config']['name']))
+                       [], self._get_layer_metadata(layer['config']['name']),
+                       self._get_layer_names_and_weights(
+                           layer['config']['name']))
 
     # There is a possible edge case where a single layer may output multiple
     # tensors and multiple tensors from that layer may be used by the
@@ -221,7 +226,8 @@ def _match_layer_with_inputs(self, layer, pattern, is_head_node):
 
     return LayerNode(layer, self._get_layer_weights(layer['config']['name']),
                      input_match_layer_nodes,
-                     self._get_layer_metadata(layer['config']['name']))
+                     self._get_layer_metadata(layer['config']['name']),
+                     self._get_layer_names_and_weights(layer['config']['name']))
 
   def _find_pattern(self, pattern, matched_layers=None):
     for layer in self._config['layers']:
@@ -265,6 +271,7 @@ def _remove_layers(self, layers_to_remove, layers_to_remove_names):
     # now that layer has been removed.
     for layer_name in layers_to_remove_names:
       self._layer_weights_map.pop(layer_name, None)
+      self._layer_names_and_weights_map.pop(layer_name, None)
       self._layer_metadata_map.pop(layer_name, None)
 
   def _replace(self, match_layer_node, replacement_layer_node):
@@ -355,8 +362,12 @@ def _add_replacement_layer(layer_node):
       """Recursively add new layers."""
       self._config['layers'].append(layer_node.layer)
       layer_name = layer_node.layer['config']['name']
+      # TODO(b/184603494): Remove weight map structure from model_transformer.
       if layer_node.weights:
         self._layer_weights_map[layer_name] = layer_node.weights
+      if layer_node.names_and_weights:
+        self._layer_names_and_weights_map[
+            layer_name] = layer_node.names_and_weights
       if layer_node.metadata:
         self._layer_metadata_map[layer_name] = layer_node.metadata
       if self.candidate_layers:
@@ -403,6 +414,9 @@ def _add_replacement_nodes(first_layer_removed_index, replacement_nodes):
         layer_name = replacement_node.layer['config']['name']
         if replacement_node.weights:
           self._layer_weights_map[layer_name] = replacement_node.weights
+        if replacement_node.names_and_weights:
+          self._layer_names_and_weights_map[
+              layer_name] = replacement_node.names_and_weights
         if replacement_node.metadata:
           self._layer_metadata_map[layer_name] = replacement_node.metadata
         if self.candidate_layers:
@@ -433,8 +447,17 @@ def _get_keras_layer_weights(self, keras_layer):
         zip(keras_layer.weights, keras_layer.get_weights()):
       weights_map[self._weight_name(weight_tensor.name)] = weight_numpy
 
+    if len(weights_map) != len(keras_layer.weights):
+      # The case that variable identifier is not unique. It's a fallback that
+      # uses weight list instead of the weights map.
+      return None
+
     return weights_map
 
+  def _get_keras_layer_names_and_weights(self, keras_layer):
+    return zip([weight.name for weight in keras_layer.weights],
+               keras_layer.get_weights())
+
   def _set_layer_weights(self, layer, weights_map):
     """Sets the values of weights in a Keras layer."""
 
@@ -447,6 +470,9 @@ def _set_layer_weights(self, layer, weights_map):
 
     K.batch_set_value(weight_value_tuples)
 
+  def _set_layer_names_and_weights(self, layer, names_and_weights):
+    layer.set_weights([weight for _, weight in names_and_weights])
+
   @staticmethod
   def _name(obj):
     return obj.__class__.__name__
@@ -496,8 +522,12 @@ def transform(self):
     # to prevent infinite loops.
     self._transform_matched_layers_map = {}
     self._layer_weights_map = {}
+    self._layer_names_and_weights_map = {}
+
     for layer in self.model.layers:
       self._layer_weights_map[layer.name] = self._get_keras_layer_weights(layer)
+      self._layer_names_and_weights_map[
+          layer.name] = self._get_keras_layer_names_and_weights(layer)
 
     # Maintains a current mutable copy of the metadata through transformation.
     self._layer_metadata_map = copy.deepcopy(self.layer_metadata)
@@ -557,5 +587,9 @@ def transform(self):
       weights = self._layer_weights_map.get(layer.name)
       if weights:
         self._set_layer_weights(layer, weights)
+      else:
+        names_and_weights = self._layer_names_and_weights_map.get(layer.name)
+        if names_and_weights:
+          self._set_layer_names_and_weights(layer, names_and_weights)
 
     return transformed_model, copy.deepcopy(self._layer_metadata_map)

tensorflow_model_optimization/python/core/quantization/keras/graph_transformations/model_transformer_test.py

@@ -75,6 +75,27 @@ def _simple_dense_model(self, model_type='functional'):
           [keras.layers.Dense(2, input_shape=(3,)),
            keras.layers.ReLU(6.0)])
 
+  def _nested_model(self, model_type='functional', submodel_type='functional'):
+    if submodel_type == 'functional':
+      inp = keras.layers.Input((3,))
+      x = keras.layers.Dense(2)(inp)
+      out = keras.layers.Dense(2)(x)
+      submodel = keras.Model(inp, out)
+    elif submodel_type == 'sequential':
+      submodel = keras.Sequential(
+          [keras.layers.Dense(2, input_shape=(3,)),
+           keras.layers.Dense(2)])
+
+    if model_type == 'functional':
+      inp = keras.layers.Input((3,))
+      x = submodel(inp)
+      out = keras.layers.ReLU(6.0)(x)
+      return keras.Model(inp, out)
+    elif model_type == 'sequential':
+      return keras.Sequential(
+          [submodel,
+           keras.layers.ReLU(6.0)])
+
   def _assert_config(self, expected_config, actual_config, exclude_keys=None):
     """Asserts that the two config dictionaries are equal.
 
@@ -681,6 +702,24 @@ def replacement(self, match_layer):
     self._assert_config(model.get_config(), transformed_model.get_config(),
                         ['build_input_shape'])
 
+  @parameterized.parameters([
+      ('sequential', 'sequential'),
+      ('sequential', 'functional'),
+      ('functional', 'sequential'),
+      ('functional', 'functional'),])
+  def testNestedModelNoChange(self, model_type, submodel_type):
+    model = self._nested_model(model_type, submodel_type)
+
+    transformed_model, _ = ModelTransformer(
+        model, []).transform()
+
+    # build_input_shape is a TensorShape object and the two objects are not
+    # considered the same even though the shapes are the same.
+    self._assert_config(model.get_config(), transformed_model.get_config(),
+                        ['class_name', 'build_input_shape'])
+
+    self._assert_model_results_equal(model, transformed_model)
+
   # Validation Tests
 
   def testRaisesErrorForSubclassModels(self):

tensorflow_model_optimization/python/core/quantization/keras/graph_transformations/transforms.py

@@ -79,26 +79,37 @@ class LayerNode(object):
   been found in a model, and layers which should be replaced inside the model.
   """
 
-  def __init__(self, layer, weights=None, input_layers=None, metadata=None):
+  def __init__(
+      self,
+      layer,
+      weights=None,
+      input_layers=None,
+      metadata=None,
+      names_and_weights=None):
     """Construct a LayerNode representing a tree of layers.
 
     Args:
       layer: layer config of this node.
       weights: An OrderedDict of weight name => value for the layer.
       input_layers: List of `LayerNode`s that feed into this layer.
       metadata: Dictionary of metadata for a given layer.
+      names_and_weights: A list of tuples (name, weight). It only used when
+        weights dictionary is empty.
     """
     if weights is None:
       weights = collections.OrderedDict()
     if input_layers is None:
       input_layers = []
     if metadata is None:
       metadata = {}
+    if names_and_weights is None:
+      names_and_weights = []
 
     self.layer = layer
     self.weights = weights
     self.input_layers = input_layers
     self.metadata = metadata
+    self.names_and_weights = names_and_weights
 
   def __str__(self):
     return '{} <- [{}]'.format(