Add remove_input_range method that removes input range after apply quantize.

PiperOrigin-RevId: 456004543

Author: Xhark

tensorflow_model_optimization/python/core/api/quantization/keras/__init__.py

@@ -30,6 +30,7 @@
 from tensorflow_model_optimization.python.core.quantization.keras.quantize import quantize_annotate_layer
 from tensorflow_model_optimization.python.core.quantization.keras.quantize import quantize_annotate_model
 from tensorflow_model_optimization.python.core.quantization.keras.quantize import quantize_apply
+from tensorflow_model_optimization.python.core.quantization.keras.quantize import remove_input_range
 
 # quantize with custom quantization parameterization or implementation, or
 # handle custom Keras layers.

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

@@ -644,3 +644,58 @@ def _is_functional_model(model):
   return (isinstance(model, keras.Model)
           and not isinstance(model, keras.Sequential)
           and model._is_graph_network)    # pylint: disable=protected-access
+
+
+def remove_input_range(model):
+  """Remove the input range.
+
+  Example:
+
+  ```python
+  model = keras.Sequential([
+      layers.Dense(10, activation='relu', input_shape=(100,)),
+      quantize_annotate_layer(layers.Dense(2, activation='sigmoid'))
+  ])
+  with quantize.quantize_scope():
+    model = quantize_annotate_model(model)
+    model = quantize_apply(model)
+    model = remove_input_range(model)
+  ```
+
+  In certain cases, a desired input range is not required if the model itself is
+  internally used.
+
+  Args:
+    model: A `tf.keras` Sequential or Functional model which has been quantized.
+
+  Returns:
+    Returns a new `tf.keras` model removed input range.
+  """
+  config = model.get_config()
+  no_input_quantizer = quantizers.NoQuantizer()
+  serialized_input_quantizer = tf.keras.utils.serialize_keras_object(
+      no_input_quantizer)
+
+  if _is_functional_model(model):
+    input_layer_list = _nested_to_flatten_node_data_list(config['input_layers'])
+    for layer_config in config['layers']:
+      input_name = _unwrap_first_input_name(layer_config['inbound_nodes'])
+      if input_name is None:
+        continue
+
+      for input_layer in input_layer_list:
+        if input_name == input_layer[0]:
+          layer_config['config']['quantizer'] = serialized_input_quantizer
+          break
+
+    model = keras.Model.from_config(config)
+  else:
+    if (len(config['layers']) < 1 or
+        config['layers'][1]['class_name'] != 'QuantizeLayer'):
+      raise ValueError('`model` should be already quantized.')
+    config['layers'][1]['config'][
+        'quantizer'] = serialized_input_quantizer
+
+    model = keras.Sequential.from_config(config)
+
+  return model

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

@@ -437,10 +437,52 @@ def __ne__(self, other):
     return not self.__eq__(other)
 
 
+class NoQuantizer(_QuantizeHelper, Quantizer):
+  """Dummy quantizer do nothing."""
+
+  def __init__(self):
+    """Construct an NoQuantizer.
+
+    This is an experimental API not subject to backward compatibility.
+    """
+    pass
+
+  def build(self, tensor_shape, name, layer):
+    pass
+
+  def __call__(self, inputs, training, weights, **kwargs):
+    """Quantize tensor.
+
+    Args:
+      inputs: Input tensor to be quantized.
+      training: Whether the graph is currently training.
+      weights: Dictionary of weights the quantizer can use to quantize the
+        tensor. This contains the weights created in the `build` function.
+      **kwargs: Additional variables which may be passed to the quantizer.
+
+    Returns:
+      Quantized tensor.
+    """
+    return inputs
+
+  def get_config(self):
+    return {}
+
+  def __eq__(self, other):
+    if not isinstance(other, NoQuantizer):
+      return False
+
+    return True
+
+  def __ne__(self, other):
+    return not self.__eq__(other)
+
+
 def _types_dict():
   return {
       'AllValuesQuantizer': AllValuesQuantizer,
       'LastValueQuantizer': LastValueQuantizer,
       'MovingAverageQuantizer': MovingAverageQuantizer,
       'FixedQuantizer': FixedQuantizer,
+      'NoQuantizer': NoQuantizer,
   }

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

@@ -36,11 +36,14 @@
     quantizers.LastValueQuantizer,
     quantizers.MovingAverageQuantizer,
     quantizers.AllValuesQuantizer,
-    quantizers.FixedQuantizer)
+    quantizers.FixedQuantizer,
+    quantizers.NoQuantizer)
 class QuantizersTest(tf.test.TestCase, parameterized.TestCase):
 
   def _get_quant_params(self, quantizer_type):
-    if quantizer_type == quantizers.FixedQuantizer:
+    if quantizer_type == quantizers.NoQuantizer:
+      return {}
+    elif quantizer_type == quantizers.FixedQuantizer:
       return {
           'num_bits': 8,
           'init_min': 0.0,