ai-edge-quantizer-nightly 0.0.1.dev20250302__py3-none-any.whl → 0.5.0.dev20260103__py3-none-any.whl

ai_edge_quantizer/utils/calibration_utils.py

@@ -15,9 +15,24 @@
 
 """Utilities for model calibration."""
 
-from typing import Union
+import copy
+from typing import Any, Union
+
 import numpy as np
+
+from ai_edge_litert.tools import flatbuffer_utils
 from ai_edge_quantizer import qtyping
+from ai_edge_quantizer.algorithms.utils import common_utils
+from ai_edge_quantizer.utils import constrained_ops_utils
+from ai_edge_quantizer.utils import tfl_flatbuffer_utils
+from ai_edge_quantizer.utils import tfl_interpreter_utils
+
+
+_SignatureInput = dict[str, Any]
+_OpQuantConstraint = common_utils.OpQuantConstraint
+_SignatureData = dict[
+    str, list[str]
+]  # signature_key -> list of signature_names.
 
 
 def _update_moving_average(
@@ -84,3 +99,250 @@ def min_max_update(qsv: qtyping.QSV, new_qsv: qtyping.QSV) -> qtyping.QSV:
   updated_qsv["min"] = np.minimum(qsv["min"], new_qsv["min"])
   updated_qsv["max"] = np.maximum(qsv["max"], new_qsv["max"])
   return updated_qsv
+
+
+def _find_overall_min_max(
+    qsv: qtyping.QSV, tensor_names: list[str]
+) -> tuple[np.ndarray, np.ndarray]:
+  """Finds the overall minimum and maximum values for the given tensors.
+
+  Args:
+    qsv: The quantization statistical value of the tensor (min/max).
+    tensor_names: The list of tensor names to find the minimum and maximum
+      values.
+
+  Returns:
+    The minimum and maximum values for the given tensors.
+  """
+  min_value = np.inf
+  max_value = -np.inf
+  for tensor_name in tensor_names:
+    min_value = min(min_value, qsv[tensor_name]["min"])
+    max_value = max(max_value, qsv[tensor_name]["max"])
+  return min_value, max_value
+
+
+class CalibrationQsvAlignmentUtils:
+  """Calibration utils for alignment of QSVs.
+
+  This class is used to align QSVs for a given model. It builds a list of ops
+  that need to be constrained to the same scale as the input. Based on this
+  list, it finds the corresponding tensor names for a given signature data.
+  """
+
+  def __init__(self, model_path: str):
+    self._same_as_input_scale_ops = (
+        constrained_ops_utils.get_constrained_op_list(
+            _OpQuantConstraint.SAME_AS_INPUT_SCALE
+        )
+    )
+
+    tfl_interpreter = tfl_interpreter_utils.create_tfl_interpreter(model_path)
+    self._flatbuffer_object = tfl_flatbuffer_utils.read_model(model_path)
+
+    signature_keys = list(tfl_interpreter.get_signature_list().keys())
+
+    # Build a dict of signature runners.
+    self._signature_runners = {}
+    for signature_key in signature_keys:
+      signature_runner = tfl_interpreter.get_signature_runner(signature_key)
+      self._signature_runners[signature_key] = signature_runner
+
+  def _search_tensor_by_signature_name(
+      self, signature_key: str, signature_input_output_name: str, verbose=False
+  ) -> list[str]:
+    """Searches for a tensor name for a given signature by signature input or output name.
+
+    Args:
+      signature_key: Name of the signature.
+      signature_input_output_name: Name of the input or output in the signature.
+      verbose: Flag to enable verbose output.
+
+    Returns:
+      The list with one or two tensor names. The first one is the input tensor
+      name, and the second one is the output tensor name.
+    """
+
+    if verbose:
+      print("Searching tensor by signature name.")
+
+    tensor_names = []
+
+    # Search among inputs.
+    input_details = self._signature_runners[signature_key].get_input_details()
+    if signature_input_output_name in input_details.keys():
+      tensor_names.append(input_details[signature_input_output_name]["name"])
+
+    # Search among outputs.
+    output_details = self._signature_runners[signature_key].get_output_details()
+    if signature_input_output_name not in output_details:
+      if not tensor_names:
+        raise ValueError(
+            f"Signature {signature_key} does not have input or output"
+            f" `{signature_input_output_name}`"
+        )
+      return tensor_names
+
+    output_tensor_name = output_details[signature_input_output_name]["name"]
+    if verbose:
+      print(
+          ">> Starting recursive search for the output tensor name:"
+          f" {output_tensor_name}"
+      )
+
+    idx = self._signature_runners[signature_key]._subgraph_index  # pylint: disable=protected-access
+    subgraph = self._flatbuffer_object.subgraphs[idx]
+    graph_info = qtyping.GraphInfo(
+        subgraph.tensors, self._flatbuffer_object.buffers
+    )
+
+    # Recursively search the graph for the output tensor name since it may be
+    # `SAME_AS_INPUT` constrainted.
+    operators = copy.deepcopy(subgraph.operators)
+    tensor_name = self._search_reverse_order_recursively(
+        graph_info, operators, output_tensor_name, indent="  ", verbose=verbose
+    )
+    tensor_names.append(tensor_name)
+
+    if verbose:
+      print(f"\n\nFound tensor name: {tensor_name}")
+
+    return tensor_names
+
+  def _search_reverse_order_recursively(
+      self,
+      graph_info: qtyping.GraphInfo,
+      operators: list[Any],
+      output_tensor_name: str,
+      indent: str,
+      verbose: bool = False,
+  ):
+    """Searches for a tensor name in reverse order recursively.
+
+    Stop criteria is when the tensor belongs to an operator that is not
+    `SAME_AS_INPUT` constrainted.
+
+    Args:
+      graph_info: Graph information.
+      operators: List of operators.
+      output_tensor_name: Name of the output tensor to search for.
+      indent: Indentation string for debug output.
+      verbose: Flag to enable verbose output.
+
+    Returns:
+      The name of the tensor found, or None if not found.
+    """
+    op_codes = self._flatbuffer_object.operatorCodes
+
+    while operators:
+      op = operators.pop()
+      op_code = op_codes[op.opcodeIndex].builtinCode
+      op_name = flatbuffer_utils.opcode_to_name(
+          self._flatbuffer_object, op.opcodeIndex
+      )
+      if op_code not in tfl_flatbuffer_utils.TFL_OP_CODE_TO_NAME:
+        continue
+      for output_idx in op.outputs:
+        if output_tensor_name == tfl_flatbuffer_utils.get_tensor_name(
+            graph_info.subgraph_tensors[output_idx]
+        ):
+          dbg_str = (
+              f"{indent}>> Found `{op_name}`, output tensor"
+              f" '{output_tensor_name}'"
+          )
+
+          if op_name not in self._same_as_input_scale_ops:
+            if verbose:
+              print(f"{dbg_str}, returning...")
+            return output_tensor_name
+
+          if verbose:
+            print(f"{dbg_str}, with SAME_AS_INPUT, search recursively among:")
+          for input_idx in op.inputs:
+            input_tensor_name = graph_info.subgraph_tensors[
+                input_idx
+            ].name.decode("utf-8")
+
+            if verbose:
+              print(f"{indent}    Input: {input_tensor_name}")
+
+            return self._search_reverse_order_recursively(
+                graph_info,
+                operators,
+                input_tensor_name,
+                indent=f"{indent}  ",
+                verbose=verbose,
+            )
+    return output_tensor_name
+
+  def align_quant_stats(
+      self, qsv: dict[str, Any], signature_data: _SignatureData
+  ) -> tuple[np.ndarray, np.ndarray]:
+    """Aligns quantization statistics for a given signature data.
+
+    This function takes quantization statistics and signature data as input,
+    identifies the tensors associated with the signature data, and aligns
+    the quantization statistics of these tensors by setting their minimum
+    and maximum values to the same value. This ensures that the tensors
+    have the same quantization parameters.
+
+    Args:
+      qsv: Quantization statistics.
+      signature_data: Signature data.
+
+    Returns:
+      Tuple of min and max values.
+    """
+    # Go over all signature info and find the corresponding tensor names.
+    tensor_names = []
+    for signature_key, signature_names in signature_data.items():
+      for signature_name in signature_names:
+        tensor_name = self._search_tensor_by_signature_name(
+            signature_key, signature_name
+        )
+        tensor_names.extend(tensor_name)
+
+    # Find min and max values accross all tensors.
+    min_value, max_value = _find_overall_min_max(qsv, tensor_names)
+
+    # Overwrite the min and max values in the QSV.
+    for tensor_name in tensor_names:
+      qsv[tensor_name]["min"] = min_value
+      qsv[tensor_name]["max"] = max_value
+
+    return min_value, max_value
+
+  def update_quant_stats(
+      self,
+      qsv: dict[str, Any],
+      signature_data: _SignatureData,
+      min_value: np.ndarray,
+      max_value: np.ndarray,
+  ):
+    """Updates quantization statistics for a given signature data.
+
+    This function updates the quantization statistics with the provided min, max
+    values for the tensors specified in the signature data.
+
+    Args:
+      qsv: Quantization statistics.
+      signature_data: Signature data.
+      min_value: Minimum value to update.
+      max_value: Maximum value to update.
+
+    Returns:
+      Updated quantization statistics.
+    """
+    # Go over all signature info and find the corresponding tensor names.
+    tensor_names = []
+    for signature_key, signature_names in signature_data.items():
+      for signature_name in signature_names:
+        tensor_name = self._search_tensor_by_signature_name(
+            signature_key, signature_name
+        )
+        tensor_names.extend(tensor_name)
+
+    # Overwrite the min and max values in the QSV.
+    for tensor_name in tensor_names:
+      qsv[tensor_name]["min"] = min_value
+      qsv[tensor_name]["max"] = max_value

ai_edge_quantizer/utils/calibration_utils_test.py

@@ -14,11 +14,67 @@
 # ==============================================================================
 
 from absl.testing import parameterized
+import numpy as np
+
 from tensorflow.python.platform import googletest
+from ai_edge_quantizer import quantizer
 from ai_edge_quantizer.utils import calibration_utils
+from ai_edge_quantizer.utils import test_utils
+from ai_edge_quantizer.utils import tfl_interpreter_utils
+
+_RNG = np.random.default_rng(66)
+
+_CALIBRATION_DATASET = {
+    "signature_1": [{
+        "cache_0": np.zeros(shape=(1, 100, 4, 4), dtype=np.float32),
+        "cache_1": np.zeros(shape=(1, 100, 4, 4), dtype=np.float32),
+        "positions": np.zeros(shape=(1, 100), dtype=np.int32),
+        "tokens": np.zeros(shape=(1, 100), dtype=np.int32),
+    }],
+    "signature_2": [{
+        "cache_0": _RNG.random(size=(1, 100, 4, 4), dtype=np.float32),
+        "cache_1": _RNG.random(size=(1, 100, 4, 4), dtype=np.float32),
+        "positions": (
+            _RNG.integers(low=0, high=10, size=(1, 100)).astype(np.int32)
+        ),
+        "tokens": _RNG.integers(low=0, high=10, size=(1, 100)).astype(np.int32),
+    }],
+}
+
 
+def _get_quant_parameters(
+    quantized_model: bytes, signature_data: dict[str, list[str]]
+) -> list[np.ndarray]:
+  """Returns the quantization parameters from the quantized model."""
+  quant_params = []
+  tfl_interpreter = tfl_interpreter_utils.create_tfl_interpreter(
+      quantized_model
+  )
+  for signature_key, signature_names in signature_data.items():
+    signature_runner = tfl_interpreter.get_signature_runner(signature_key)
+
+    for signature_name in signature_names:
+      input_details = signature_runner.get_input_details()
+      output_details = signature_runner.get_output_details()
+      if signature_name in input_details.keys():
+        quant_param = input_details[signature_name]["quantization_parameters"][
+            "scales"
+        ].squeeze()
+        quant_params.append(quant_param)
+      elif signature_name in output_details.keys():
+        output_details = signature_runner.get_output_details()
+        quant_param = output_details[signature_name]["quantization_parameters"][
+            "scales"
+        ].squeeze()
+        quant_params.append(quant_param)
+      else:
+        raise ValueError(
+            f"Signature name {signature_name} not found in the model."
+        )
+  return quant_params
 
-class CalibrationUtilsTest(parameterized.TestCase):
+
+class CalibrationQsvAlignmentUtilsTest(parameterized.TestCase):
 
   @parameterized.named_parameters(
       dict(
@@ -66,12 +122,126 @@ class CalibrationUtilsTest(parameterized.TestCase):
   def test_update_tensor_qsv_min_max(self, old_qsv, new_qsv, expected_qsv):
     updated_qsv = calibration_utils.min_max_update(old_qsv, new_qsv)
     if isinstance(expected_qsv["min"], list):
-      self.assertListEqual(list(updated_qsv["min"]), expected_qsv["min"])
-      self.assertListEqual(list(updated_qsv["max"]), expected_qsv["max"])
+      self.assertEqual(list(updated_qsv["min"]), expected_qsv["min"])
+      self.assertEqual(list(updated_qsv["max"]), expected_qsv["max"])
     else:
       self.assertEqual(updated_qsv["min"], expected_qsv["min"])
       self.assertEqual(updated_qsv["max"], expected_qsv["max"])
 
+  def test_calibration_utils_init_fails(self):
+    model_path = "non_existent_model.tflite"
+    with self.assertRaisesWithPredicateMatch(
+        Exception, lambda err: f"{model_path}" in str(err)
+    ):
+      calibration_utils.CalibrationQsvAlignmentUtils(model_path)
+
+  def test_calibration_utils_init_succeeds(self):
+    model_path = test_utils.get_path_to_datafile(
+        "../tests/models/single_add.tflite"
+    )
+    calib_utils = calibration_utils.CalibrationQsvAlignmentUtils(model_path)
+    self.assertNotEmpty(calib_utils._signature_runners)
+    self.assertNotEmpty(calib_utils._same_as_input_scale_ops)
+
+  def test_search_tensor_by_signature_name_succeeds_on_unconstrained_op(self):
+    model_path = test_utils.get_path_to_datafile(
+        "../tests/models/single_add.tflite"
+    )
+    expected_tensor_name = "PartitionedCall:0"
+    calib_utils = calibration_utils.CalibrationQsvAlignmentUtils(model_path)
+    tensor_name = calib_utils._search_tensor_by_signature_name(
+        "serving_default", "add"
+    )
+    self.assertEqual(tensor_name, [expected_tensor_name])
+
+  def test_search_tensor_by_signature_name_succeeds_on_constrained_op(self):
+    model_path = test_utils.get_path_to_datafile(
+        "../tests/models/single_slice.tflite"
+    )
+    expected_tensor_name = "slice_input_tensor:0"
+    calib_utils = calibration_utils.CalibrationQsvAlignmentUtils(model_path)
+    tensor_name = calib_utils._search_tensor_by_signature_name(
+        "slice", "output_0"
+    )
+    self.assertEqual(tensor_name, [expected_tensor_name])
+
+  def test_align_quant_stats_succeeds(self):
+    model_path = test_utils.get_path_to_datafile(
+        "../tests/models/toy_model_with_kv_cache_multi_signature.tflite"
+    )
+    recipe_path = test_utils.get_path_to_datafile(
+        "../recipes/default_a8w8_recipe.json"
+    )
+    signature_data = {
+        "signature_1": ["output_1_1"],
+        "signature_2": ["output_1_1"],
+    }
+
+    # Obtain the calibration results.
+    qt = quantizer.Quantizer(model_path, recipe_path)
+    qsv = qt.calibrate(_CALIBRATION_DATASET)
+
+    # First quantize the model without aligning the quantization parameters.
+    quantized_model = qt.quantize(qsv).quantized_model
+    quant_params = _get_quant_parameters(quantized_model, signature_data)
+    self.assertFalse(
+        all(x == quant_params[0] for x in quant_params)
+    )  # not equal quantization params.
+
+    # Align the quantization parameters and quantize again.
+    calib_utils = calibration_utils.CalibrationQsvAlignmentUtils(model_path)
+    calib_utils.align_quant_stats(qsv, signature_data)
+    quantized_model = qt.quantize(qsv).quantized_model
+    quant_params = _get_quant_parameters(quantized_model, signature_data)
+    self.assertTrue(
+        all(x == quant_params[0] for x in quant_params)
+    )  # equal quantization params.
+
+  def test_update_quant_stats_succeeds(self):
+    model_path = test_utils.get_path_to_datafile(
+        "../tests/models/toy_model_with_kv_cache_multi_signature.tflite"
+    )
+    recipe_path = test_utils.get_path_to_datafile(
+        "../recipes/default_a8w8_recipe.json"
+    )
+    signature_data = {
+        "signature_1": ["output_1_1"],
+        "signature_2": ["output_1_1"],
+    }
+
+    # Obtain the calibration results.
+    qt = quantizer.Quantizer(model_path, recipe_path)
+    qsv = qt.calibrate(_CALIBRATION_DATASET)
+
+    # First quantize the model without updating the `signature_1`.
+    quantized_model = qt.quantize(qsv).quantized_model
+    quant_params = _get_quant_parameters(quantized_model, signature_data)
+    self.assertFalse(
+        all(x == quant_params[0] for x in quant_params)
+    )  # not equal quantization params.
+
+    # Update the `signature_1` with stats from `signature_2`.
+    calib_utils = calibration_utils.CalibrationQsvAlignmentUtils(model_path)
+    min_val, max_val = calib_utils.align_quant_stats(  # for min and max only.
+        qsv,
+        {
+            "signature_2": ["output_1_1"],
+        },
+    )
+    calib_utils.update_quant_stats(
+        qsv,
+        {
+            "signature_1": ["output_1_1"],
+        },
+        min_val,
+        max_val,
+    )
+    quantized_model = qt.quantize(qsv).quantized_model
+    quant_params = _get_quant_parameters(quantized_model, signature_data)
+    self.assertTrue(
+        all(x == quant_params[0] for x in quant_params)
+    )  # equal quantization params.
+
 
 if __name__ == "__main__":
   googletest.main()

ai_edge_quantizer/utils/constrained_ops_utils.py

@@ -0,0 +1,111 @@
+# Copyright 2024 The AI Edge Quantizer Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+"""Utils for handling operators with quantization constraints."""
+
+from ai_edge_quantizer import algorithm_manager
+from ai_edge_quantizer import qtyping
+from ai_edge_quantizer.algorithms.uniform_quantize import common_quantize
+from ai_edge_quantizer.algorithms.utils import common_utils
+from ai_edge_litert import schema_py_generated as schema_fb  # pylint: disable=g-direct-tensorflow-import
+
+
+_OpQuantConstraint = common_utils.OpQuantConstraint
+
+
+def get_constrained_op_list(
+    quant_constraint: _OpQuantConstraint, verbose: bool = False
+) -> list[str]:
+  """Constructs and returns a list of constrained operators.
+
+  This is achieved by invoking all materialization functions and extracting
+  the constraint argument, using monkey patching to redirect logic to wrapper
+  functions.
+
+  Args:
+    quant_constraint: The quantization constraint to filter operators by.
+    verbose: Flag to enable verbose output.
+
+  Returns:
+    A list containing operators with the specified constraint.
+  """
+  constrained_ops = []
+
+  def materialize_standard_op_wrapper(
+      op_info: qtyping.OpInfo,
+      *_args,
+      constraint: _OpQuantConstraint = _OpQuantConstraint.NO_CONSTRAIN,
+      **_kwargs,
+  ) -> list[qtyping.TensorTransformationParams]:
+    if constraint == quant_constraint:
+      constrained_ops.append(op_info.op_name)
+    # Return dummy values to avoid exceptions.
+    dummy_value = [qtyping.TensorTransformationParams("")] * 2
+    return dummy_value
+
+  # Dummy implementation of the `_are_weights_too_small` function to support
+  # `materialize_standard_op_wrapper` above.
+  def are_weights_too_small_wrapper(*_args, **_kwargs) -> bool:
+    return False
+
+  # Dummy implementation of the `_materialize_bias_for_fc_conv_ops` function to
+  # support `materialize_standard_op_wrapper` above.
+  def materialize_bias_for_fc_conv_ops_wrapper(*_args, **_kwargs):
+    return
+
+  # Do monkey patch to intercept the `materialize_standard_op` function to
+  # support `materialize_standard_op_wrapper` above.
+  original_materialize_standard_op = common_utils.materialize_standard_op
+  original_are_weights_too_small = common_quantize._are_weights_too_small  # pylint: disable=protected-access
+  original_materialize_bias_for_fc_conv_ops = (
+      common_quantize._materialize_bias_for_fc_conv_ops  # pylint: disable=protected-access
+  )
+  common_utils.materialize_standard_op = materialize_standard_op_wrapper
+  common_quantize._are_weights_too_small = are_weights_too_small_wrapper  # pylint: disable=protected-access
+  common_quantize._materialize_bias_for_fc_conv_ops = (  # pylint: disable=protected-access
+      materialize_bias_for_fc_conv_ops_wrapper
+  )
+  minmax_func_dict = algorithm_manager.MIN_MAX_OP_NAME_MATERIALIZE_FUNC_DICT
+
+  # Loop over all available materialization functions to build up a list of
+  # ops with the given constraint.
+  for op, materialize_fn in minmax_func_dict.items():
+    # Create a dummy op info to trigger the materialization.
+    mock_op = schema_fb.OperatorT()
+    mock_op.inputs = [0]
+    mock_op.outputs = [0]
+    op_info = qtyping.OpInfo(
+        op=mock_op,
+        op_name=op,
+        subgraph_op_index=0,
+        op_quant_config=qtyping.OpQuantizationConfig(),
+    )
+    materialize_fn(
+        get_tensor_quant_params_fn=None,
+        op_info=op_info,
+        graph_info=None,
+        tensor_name_to_qsv=None,
+    )
+
+  if verbose:
+    print(f"  {quant_constraint} op list: {constrained_ops}")
+
+  # Restore the original functions.
+  common_utils.materialize_standard_op = original_materialize_standard_op
+  common_quantize._are_weights_too_small = original_are_weights_too_small  # pylint: disable=protected-access
+  common_quantize._materialize_bias_for_fc_conv_ops = (  # pylint: disable=protected-access
+      original_materialize_bias_for_fc_conv_ops
+  )
+  return constrained_ops

ai_edge_quantizer/utils/constrained_ops_utils_test.py

@@ -0,0 +1,50 @@
+# Copyright 2024 The AI Edge Quantizer Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+from tensorflow.python.platform import googletest
+from absl.testing import parameterized
+from ai_edge_quantizer.algorithms.utils import common_utils
+from ai_edge_quantizer.utils import constrained_ops_utils
+
+
+_OpQuantConstraint = common_utils.OpQuantConstraint
+
+
+class ConstrainedOpsUtilsTest(parameterized.TestCase):
+
+  @parameterized.named_parameters(
+      dict(
+          testcase_name="same_as_input_scale",
+          constraint=_OpQuantConstraint.SAME_AS_INPUT_SCALE,
+          expected_num_ops=18,
+      ),
+      dict(
+          testcase_name="same_as_output_scale",
+          constraint=_OpQuantConstraint.SAME_AS_OUTPUT_SCALE,
+          expected_num_ops=7,
+      ),
+      dict(
+          testcase_name="no_constrain",
+          constraint=_OpQuantConstraint.NO_CONSTRAIN,
+          expected_num_ops=25,
+      ),
+  )
+  def test_get_constrained_op_list(self, constraint, expected_num_ops):
+    constrained_ops = constrained_ops_utils.get_constrained_op_list(constraint)
+    self.assertLen(constrained_ops, expected_num_ops)
+
+
+if __name__ == "__main__":
+  googletest.main()