Added an argument to support user-provided name prefix

PiperOrigin-RevId: 491514531

Author: tensorflower-gardener

tensorflow_model_optimization/python/core/quantization/keras/quantize.py

@@ -80,7 +80,7 @@ def quantize_scope(*args):
   return tf.keras.utils.custom_object_scope(*(args + (quantization_objects,)))
 
 
-def quantize_model(to_quantize):
+def quantize_model(to_quantize, quantized_layer_name_prefix='quant_'):
   """Quantize a `tf.keras` model with the default quantization implementation.
 
   Quantization constructs a model which emulates quantization during training.
@@ -117,13 +117,18 @@ def quantize_model(to_quantize):
   Args:
     to_quantize: tf.keras model to be quantized. It can have pre-trained
       weights.
+    quantized_layer_name_prefix: Name prefix for the quantized layers. The
+      default is `quant_`.
 
   Returns:
     Returns a new `tf.keras` model prepared for quantization.
   """
   if to_quantize is None:
     raise ValueError('`to_quantize` cannot be None')
 
+  if quantized_layer_name_prefix is None:
+    quantized_layer_name_prefix = ''
+
   if not isinstance(to_quantize, keras.Model):
     raise ValueError(
         '`to_quantize` can only be a `tf.keras.Model` instance. Use '
@@ -138,7 +143,8 @@ def quantize_model(to_quantize):
         'Functional model.')
 
   annotated_model = quantize_annotate_model(to_quantize)
-  return quantize_apply(annotated_model)
+  return quantize_apply(
+      annotated_model, quantized_layer_name_prefix=quantized_layer_name_prefix)
 
 
 def quantize_annotate_model(to_annotate):
@@ -281,7 +287,8 @@ def quantize_annotate_layer(to_annotate, quantize_config=None):
 @metrics.MonitorBoolGauge('quantize_apply_usage')
 def quantize_apply(
     model,
-    scheme=default_8bit_quantize_scheme.Default8BitQuantizeScheme()):
+    scheme=default_8bit_quantize_scheme.Default8BitQuantizeScheme(),
+    quantized_layer_name_prefix='quant_'):
   """Quantize a `tf.keras` model that has been annotated for quantization.
 
   Quantization constructs a model which emulates quantization during training.
@@ -319,6 +326,8 @@ def quantize_apply(
       with `quantize_annotate`. It can have pre-trained weights.
     scheme: A `QuantizeScheme` which specifies transformer and quantization
       registry. The default is `Default8BitQuantizeScheme()`.
+    quantized_layer_name_prefix: A name prefix for quantized layers. The default
+      is `quant_`.
 
   Returns:
     Returns a new `tf.keras` model in which the annotated layers have been
@@ -327,6 +336,9 @@ def quantize_apply(
   if model is None:
     raise ValueError('`model` cannot be None')
 
+  if quantized_layer_name_prefix is None:
+    quantized_layer_name_prefix = ''
+
   if not isinstance(model, keras.Model):
     raise ValueError('`model` can only be a `tf.keras.Model` instance.'
                      'You passed an instance of type: {input}.'.format(
@@ -435,7 +447,7 @@ def _quantize(layer):  # pylint: disable=missing-docstring
     # `QuantizeAnnotate` wrapper may contain `batch_input_shape` like params.
     # TODO(pulkitb): Ensure this does not affect model cloning.
     return quantize_wrapper.QuantizeWrapperV2(
-        layer, quantize_config)
+        layer, quantize_config, name_prefix=quantized_layer_name_prefix)
 
   # 1. Create a copy of the model with the same weights. This ensures
   # modifications don't affect the original model, or its weights.
@@ -446,7 +458,7 @@ def _quantize(layer):  # pylint: disable=missing-docstring
         'Unable to clone model. This generally happens if you used custom '
         'Keras layers or objects in your model. Please specify them via '
         '`quantize_scope` for your calls to `quantize_model` and '
-        '`quantize_apply`. [%s].' % er)
+        '`quantize_apply`. [%s].' % er) from er
 
   # 2. Remove QuantizeAnnotate wrappers from the layers in the model. This
   # extracts the original model structure (easier to transform), and

tensorflow_model_optimization/python/core/quantization/keras/quantize_models_test.py

@@ -22,7 +22,6 @@
 import tempfile
 
 from absl.testing import parameterized
-
 import numpy as np
 import tensorflow as tf
 
@@ -108,13 +107,55 @@ def testModelEndToEnd(self, model_type):
     model.fit(x_train, y_train)
 
     # 3. Ensure conversion to TFLite works.
-    _, tflite_file = tempfile.mkstemp('.tflite')
-    print('TFLite File: ', tflite_file)
-    with quantize.quantize_scope():
-      utils.convert_keras_to_tflite(model, tflite_file)
+    with tempfile.NamedTemporaryFile(suffix='.tflite') as t:
+      with quantize.quantize_scope():
+        utils.convert_keras_to_tflite(model, t.name)
+
+      # 4. Verify input runs on converted model.
+      self._verify_tflite(t.name, x_train, y_train)
+
+  # Test the model with custom layer name prefix.
+  @parameterized.product(
+      model_type=_KERAS_APPLICATION_MODELS,
+      name_prefix=['', 'custom_prefix_'])
+  def testModelEndToEndCustomNamePrefix(self, model_type, name_prefix):
+    # 1. Check whether quantized model graph can be constructed.
+    model = self._get_model(model_type)
+    original_layer_names = set([layer.name for layer in model.layers])
+
+    model = quantize.quantize_model(
+        model, quantized_layer_name_prefix=name_prefix)
+    quantized_layer_names = set([layer.name for layer in model.layers])
+
+    # Remove the name of layer which is newly added to quantize the input.
+    quantized_layer_names.remove('quantize_layer')
+
+    if not name_prefix or name_prefix is None:
+      # The set of layer names should be the same.
+      self.assertEqual(original_layer_names, quantized_layer_names)
+    else:
+      self.assertNotEqual(original_layer_names, quantized_layer_names)
+      for name in original_layer_names:
+        if name in quantized_layer_names:
+          quantized_layer_names.remove(name)
+        elif name_prefix + name in quantized_layer_names:
+          quantized_layer_names.remove(name_prefix + name)
+
+      self.assertEmpty(quantized_layer_names)
+
+    # 2. Sanity check to ensure basic training on random data works.
+    x_train, y_train = self._create_test_data(model)
+    model.compile(
+        loss='categorical_crossentropy', optimizer='sgd', metrics=['accuracy'])
+    model.fit(x_train, y_train)
+
+    # 3. Ensure conversion to TFLite works.
+    with tempfile.NamedTemporaryFile(suffix='.tflite') as t:
+      with quantize.quantize_scope():
+        utils.convert_keras_to_tflite(model, t.name)
 
-    # 4. Verify input runs on converted model.
-    self._verify_tflite(tflite_file, x_train, y_train)
+      # 4. Verify input runs on converted model.
+      self._verify_tflite(t.name, x_train, y_train)
 
 
 if __name__ == '__main__':

tensorflow_model_optimization/python/core/quantization/keras/quantize_wrapper.py

@@ -41,17 +41,22 @@
 class QuantizeWrapper(tf.keras.layers.Wrapper):
   """Quantizes the weights and activations of the keras layer it wraps."""
 
-  def __init__(self, layer, quantize_config, **kwargs):
+  def __init__(self, layer, quantize_config, name_prefix='quant_', **kwargs):
     """Create a quantize emulate wrapper for a keras layer.
 
     Args:
       layer: The keras layer to be quantized.
       quantize_config: `QuantizeConfig` to quantize layer.
+      name_prefix: Prefix for quantized keras layer name. The default is
+        `quant_`.
       **kwargs: Additional keyword arguments to be passed to the keras layer.
     """
     if layer is None:
       raise ValueError('`layer` cannot be None.')
 
+    if name_prefix is None:
+      name_prefix = ''
+
     # Check against keras.Model since it is an instance of keras.layers.Layer.
     if not isinstance(layer, tf.keras.layers.Layer) or isinstance(
         layer, tf.keras.Model):
@@ -65,7 +70,7 @@ def __init__(self, layer, quantize_config, **kwargs):
                        'quantize a layer.')
 
     if 'name' not in kwargs:
-      kwargs['name'] = self._make_layer_name(layer)
+      kwargs['name'] = self._make_layer_name(layer, name_prefix)
 
     super(QuantizeWrapper, self).__init__(layer, **kwargs)
     self.quantize_config = quantize_config
@@ -74,8 +79,8 @@ def __init__(self, layer, quantize_config, **kwargs):
     metrics.MonitorBoolGauge('quantize_wrapper_usage').set(
         layer.__class__.__name__)
 
-  def _make_layer_name(self, layer):
-    return '{}_{}'.format('quant', layer.name)
+  def _make_layer_name(self, layer, name_prefix):
+    return '{}{}'.format(name_prefix, layer.name)
 
   def _weight_name(self, name):
     """Extracts the weight name from the full TensorFlow variable name.
@@ -100,8 +105,8 @@ def build(self, input_shape):
         trainable=False)
 
     self._weight_vars = []
-    for weight, quantizer in \
-        self.quantize_config.get_weights_and_quantizers(self.layer):
+    for weight, quantizer in (
+        self.quantize_config.get_weights_and_quantizers(self.layer)):
       quantizer_vars = quantizer.build(weight.shape,
                                        self._weight_name(weight.name), self)
 
@@ -110,8 +115,8 @@ def build(self, input_shape):
       self._trainable_weights.append(weight)
 
     self._quantize_activations = []
-    for activation, quantizer in \
-        self.quantize_config.get_activations_and_quantizers(self.layer):
+    for activation, quantizer in (
+        self.quantize_config.get_activations_and_quantizers(self.layer)):
       quantize_activation = quantize_aware_activation.QuantizeAwareActivation(
           activation, quantizer, self.optimizer_step, self)