Add fix_input_output_range method that fix the input and output range after apply quantize.

PiperOrigin-RevId: 427484618

Author: Xhark

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

@@ -26,6 +26,7 @@
 from tensorflow_model_optimization.python.core.quantization.keras.quantize import quantize_model
 
 # quantize some layers with default or custom quantization implementation.
+from tensorflow_model_optimization.python.core.quantization.keras.quantize import fix_input_output_range
 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

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

@@ -17,6 +17,7 @@
 # quantize with custom quantization parameterization or implementation, or
 # handle custom Keras layers.
 from tensorflow_model_optimization.python.core.quantization.keras.quantizers import AllValuesQuantizer
+from tensorflow_model_optimization.python.core.quantization.keras.quantizers import FixedQuantizer
 from tensorflow_model_optimization.python.core.quantization.keras.quantizers import LastValueQuantizer
 from tensorflow_model_optimization.python.core.quantization.keras.quantizers import MovingAverageQuantizer
 from tensorflow_model_optimization.python.core.quantization.keras.quantizers import Quantizer

tensorflow_model_optimization/python/core/quantization/keras/BUILD

@@ -90,6 +90,8 @@ py_strict_library(
     visibility = ["//visibility:public"],
     deps = [
         # six dep1,
+        # tensorflow dep1,
+        "//tensorflow_model_optimization/python/core/quantization/keras:quantizers",
     ],
 )
 
@@ -270,6 +272,7 @@ py_strict_test(
         ":quantize_layer",
         ":quantize_wrapper",
         ":quantizers",
+        # absl/testing:parameterized dep1,
         # numpy dep1,
         # tensorflow dep1,
         "//tensorflow_model_optimization/python/core/keras:test_utils",

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

@@ -25,14 +25,17 @@
 from tensorflow_model_optimization.python.core.keras import compat as tf_compat
 
 
-def FixedQuantize(inputs, init_min=-6.0, init_max=6.0, scope=None):
+def FixedQuantize(
+    inputs, init_min=-6.0, init_max=6.0, scope=None, narrow_range=False):
   """Adds a fake quantize layer with fixed quantization interval.
 
   Args:
     inputs: a tensor containing values to be quantized.
     init_min: the lower end of quantization interval.
     init_max: the upper end of quantization interval.
     scope: Optional scope for name_scope.
+    narrow_range: Whether to use the narrow quantization range
+      [1; 2^num_bits - 1] or wide range [0; 2^num_bits - 1].
   Returns:
     a tensor containing quantized values.
   """
@@ -41,7 +44,7 @@ def FixedQuantize(inputs, init_min=-6.0, init_max=6.0, scope=None):
 
   with tf.name_scope(scope):
     return tf.quantization.fake_quant_with_min_max_args(
-        inputs, min=init_min, max=init_max)
+        inputs, min=init_min, max=init_max, narrow_range=narrow_range)
 
 
 def AllValuesQuantize(inputs,

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

@@ -71,6 +71,7 @@ def quantize_scope(*args):
       'QuantizeWrapperV2': quantize_wrapper.QuantizeWrapperV2,
       'QuantizeLayer': quantize_layer.QuantizeLayer,
       'OutputOnlyConfig': quantize_config_mod.OutputOnlyConfig,
+      'FixedQuantizeConfig': quantize_config_mod.FixedQuantizeConfig,
   }
   quantization_objects.update(default_8bit_quantize_registry._types_dict())  # pylint: disable=protected-access
   quantization_objects.update(default_n_bit_quantize_registry._types_dict())  # pylint: disable=protected-access
@@ -472,3 +473,169 @@ def _quantize(layer):  # pylint: disable=missing-docstring
 
   return keras.models.clone_model(
       transformed_model, input_tensors=None, clone_function=_quantize)
+
+
+def _unwrap_first_input_name(inbound_nodes):
+  """Unwrap inbound_nodes three times to get first input name.
+
+  Args:
+    inbound_nodes: A str config that indicates input node. This method assumed
+      the inbound_nodes looks like `[[['input', 0, 0, {}]]]`.
+
+  Returns:
+    Returns a str name for the first inbound node.
+  """
+  current = inbound_nodes
+
+  for _ in range(3):
+    if not current:
+      return None
+    if not isinstance(current, list):
+      return None
+    current = current[0]
+
+  if isinstance(current, str):
+    return current
+
+  return None
+
+
+def _wrap_fixed_range(
+    quantize_config, num_bits, init_min, init_max, narrow_range):
+  config = quantize_config_mod.FixedQuantizeConfig.from_config(
+      {'config': quantize_config,
+       'num_bits': num_bits,
+       'init_min': init_min,
+       'init_max': init_max,
+       'narrow_range': narrow_range})
+  return tf.keras.utils.serialize_keras_object(config)
+
+
+def _is_serialized_node_data(nested):
+  # Node data can be of form `[layer_name, node_id, tensor_id]` or
+  # `[layer_name, node_id, tensor_id, kwargs]`.
+  if (isinstance(nested, list) and (len(nested) in [3, 4]) and
+      isinstance(nested[0], str)):
+    return True
+  return False
+
+
+def _nested_to_flatten_node_data_list(nested):
+  """Makes nested node data to flatten node data list."""
+  if _is_serialized_node_data(nested):
+    return [nested]
+
+  if isinstance(nested, list):
+    return sum(map(_nested_to_flatten_node_data_list, nested), [])
+
+  if isinstance(nested, dict):
+    return sum(map(_nested_to_flatten_node_data_list, nested.values()), [])
+
+  raise ValueError('{} is not a supported nested node data.'.format(nested))
+
+
+def fix_input_output_range(
+    model,
+    num_bits=8,
+    input_min=0.0,
+    input_max=1.0,
+    output_min=0.0,
+    output_max=1.0,
+    narrow_range=False):
+  """Fix the input and output ranges.
+
+  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 = fix_input_output_range(model, num_bits=4,
+        input_min=0, input_max=15,
+        output_min=0, output_max=15,
+        narrow_range=False)
+  ```
+
+  In certain cases, a desired input/output ranges is known and should not be
+  altered during training. To set these values, use the arguments as follows:
+
+  Args:
+    model: A `tf.keras` Sequential or Functional model which has been quantized.
+    num_bits: Number of bits for quantization
+    input_min: The lower end of quantization interval for the input.
+    input_max: The upper end of quantization interval for the input.
+    output_min: The lower end of quantization interval for the output.
+    output_max: The upper end of quantization interval for the output.
+    narrow_range: In case of 8 bits, narrow_range nudges the quantized range
+      to be [-127, 127] instead of [-128, 127]. This ensures symmetric
+      range has 0 as the centre.
+
+  Returns:
+    Returns a new `tf.keras` model fixed input range set to (input_min,
+    input_max) and fixed output range set to (output_min, output_max).
+  """
+  config = model.get_config()
+  fixed_input_quantizer = quantizers.FixedQuantizer(
+      num_bits=num_bits,
+      init_min=input_min,
+      init_max=input_max,
+      narrow_range=narrow_range)
+  serialized_fixed_input_quantizer = tf.keras.utils.serialize_keras_object(
+      fixed_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_fixed_input_quantizer
+          break
+
+    output_layer_list = _nested_to_flatten_node_data_list(
+        config['output_layers'])
+    for layer_config in config['layers']:
+      for output_layer in output_layer_list:
+        if layer_config['config']['name'] == output_layer[0]:
+          if 'quantize_config' in layer_config['config']:
+            layer_config['config']['quantize_config'] = (
+                _wrap_fixed_range(
+                    layer_config['config']['quantize_config'],
+                    num_bits=num_bits,
+                    init_min=output_min,
+                    init_max=output_max,
+                    narrow_range=narrow_range))
+          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_fixed_input_quantizer
+    if 'quantize_config' in config['layers'][-1]['config']:
+      config['layers'][-1]['config']['quantize_config'] = (
+          _wrap_fixed_range(
+              config['layers'][-1]['config']['quantize_config'],
+              num_bits=num_bits,
+              init_min=output_min,
+              init_max=output_max,
+              narrow_range=narrow_range))
+
+    model = keras.Sequential.from_config(config)
+
+  return model
+
+
+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

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

@@ -17,6 +17,9 @@
 import abc
 import six
 
+import tensorflow as tf
+from tensorflow_model_optimization.python.core.quantization.keras import quantizers
+
 
 @six.add_metaclass(abc.ABCMeta)
 class QuantizeConfig(object):
@@ -215,3 +218,54 @@ def get_config(self):
   @classmethod
   def from_config(cls, config):
     return cls(**config)
+
+
+class FixedQuantizeConfig(QuantizeConfig):
+  """QuantizeConfig that quantizes output with fixed range."""
+
+  def __init__(self, config, num_bits, init_min, init_max, narrow_range):
+    self.config = config
+    self.num_bits = num_bits
+    self.init_min = init_min
+    self.init_max = init_max
+    self.narrow_range = narrow_range
+    self.fixed_quantizer = quantizers.FixedQuantizer(
+        num_bits=num_bits,
+        init_min=init_min,
+        init_max=init_max,
+        narrow_range=narrow_range)
+
+  def get_weights_and_quantizers(self, layer):
+    return self.config.get_weights_and_quantizers(layer)
+
+  def set_quantize_weights(self, layer, quantize_weights):
+    return self.config.set_quantize_weights(layer, quantize_weights)
+
+  def get_activations_and_quantizers(self, layer):
+    activations_and_quantizers = (
+        self.config.get_activations_and_quantizers(layer))
+    return [(activation, self.fixed_quantizer)
+            for activation, _ in activations_and_quantizers]
+
+  def set_quantize_activations(self, layer, quantize_activations):
+    return self.config.set_quantize_activations(
+        layer, quantize_activations)
+
+  def get_output_quantizers(self, layer):
+    outputs_and_quantizers = (
+        self.config.get_output_quantizers(layer))
+    return [self.fixed_quantizer
+            for _ in outputs_and_quantizers]
+
+  def get_config(self):
+    return {
+        'config': tf.keras.utils.serialize_keras_object(self.config),
+        'num_bits': self.num_bits,
+        'init_min': self.init_min,
+        'init_max': self.init_max,
+        'narrow_range': self.narrow_range}
+
+  @classmethod
+  def from_config(cls, config):
+    config['config'] = tf.keras.utils.deserialize_keras_object(config['config'])
+    return cls(**config)