Only set weights returned by Transform.

Currently, ModelTransformer expects all the weights
required by a layer to be returned in the Transform
and set. This has issues since the Transform might
not know all the weights, or might only want to set
some of the weights, not all. Some of the weights
are often created by the layer in the build method.

Example: ConvBatchNorm layers create quant variables
which are determined by the weights and quantizers.
The Transform does not know about them.

This CL allows the Transform to set specific weights
it wants. Sets only the map values, not the entire list.

PiperOrigin-RevId: 278967295

Author: nutsiepully

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

@@ -19,6 +19,7 @@
 import copy
 
 from tensorflow.python import keras
+from tensorflow.python.keras import backend as K
 
 from tensorflow_model_optimization.python.core.quantization.keras.graph_transformations import transforms as transforms_mod
 
@@ -302,15 +303,41 @@ def _add_replacement_layer(layer_node):
 
     _add_replacement_layer(replacement_layer_node)
 
+  @staticmethod
+  def _weight_name(name):
+    """Extracts the weight name by removing layer from TF variable name.
+
+    For example, returns 'kernel:0' for 'dense_2/kernel:0'.
+
+    Args:
+      name: TensorFlow variable name.
+
+    Returns:
+      Extracted weight name.
+    """
+    return name.split('/')[-1]
+
   def _get_keras_layer_weights(self, keras_layer):
     """Returns a map of weight name, weight matrix. Keeps keras ordering."""
     weights_map = collections.OrderedDict()
     for weight_tensor, weight_numpy in \
         zip(keras_layer.weights, keras_layer.get_weights()):
-      weights_map[weight_tensor.name] = weight_numpy
+      weights_map[self._weight_name(weight_tensor.name)] = weight_numpy
 
     return weights_map
 
+  def _set_layer_weights(self, layer, weights_map):
+    """Sets the values of weights in a Keras layer."""
+
+    weight_value_tuples = []
+    for weight_tensor in layer.weights:
+      weight_name = self._weight_name(weight_tensor.name)
+      if weight_name in weights_map:
+        weight_value_tuples.append(
+            (weight_tensor, weights_map[weight_name]))
+
+    K.batch_set_value(weight_value_tuples)
+
   def transform(self):
     """Transforms the Keras model by applying all the specified transforms.
 
@@ -390,7 +417,7 @@ def transform(self):
     for layer in transformed_model.layers:
       weights = self._layer_weights_map.get(layer.name)
       if weights:
-        layer.set_weights(list(weights.values()))
+        self._set_layer_weights(layer, weights)
 
     # TODO(pulkitb): Consider returning the updated metadata for the
     # transformed model along with the model. This allows the opportunity for

tensorflow_model_optimization/python/core/quantization/keras/tflite/tflite_transforms.py

@@ -37,27 +37,10 @@ def _get_conv_bn_layers(bn_layer_node):
 
 def _get_weights(bn_layer_node):
   """Returns weight values for fused layer, including copying original values in unfused version."""
-  weights = collections.OrderedDict()
-
-  bn_layer_weights = list(bn_layer_node.weights.items())
-  conv_layer_weights = list(bn_layer_node.input_layers[0].weights.items())
-
-  weights['conv/kernel'] = conv_layer_weights[0][1]
-  weights['batch_normalization/gamma:0'] = bn_layer_weights[0][1]
-  weights['batch_normalization/beta:0'] = bn_layer_weights[1][1]
-
-  # TODO(tfmot): remove hardcoded initialization values.
-  weights['weight_min'] = np.array(-6.0)
-  weights['weight_max'] = np.array(6.0)
-  weights['optimizer_step'] = np.array(-1)
-  weights['activation_min'] = np.array(-6.0)
-  weights['activation_max'] = np.array(6.0)
-
-  weights['batch_normalization/moving_mean:0'] = bn_layer_weights[2][1]
-  weights['batch_normalization/moving_variance:0'] = bn_layer_weights[3][1]
-
-  return weights
 
+  return collections.OrderedDict(
+      list(bn_layer_node.input_layers[0].weights.items())
+      + list(bn_layer_node.weights.items()))
 
 def _get_params(conv_layer, bn_layer, relu_layer=None):
   """Retrieve conv_bn params within wrapped layers."""