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

ai_edge_quantizer/utils/test_utils.py

@@ -18,12 +18,55 @@
 import inspect as _inspect
 import os.path as _os_path
 import sys as _sys
-from typing import Any, Union
+from typing import Optional, Union
 
-import numpy as np
+from absl.testing import parameterized
 
+from ai_edge_quantizer import model_validator
+from ai_edge_quantizer import qtyping
+from ai_edge_quantizer import quantizer
 from ai_edge_quantizer.utils import tfl_interpreter_utils
 
+_ComputePrecision = qtyping.ComputePrecision
+_OpName = qtyping.TFLOperationName
+_TensorQuantConfig = qtyping.TensorQuantizationConfig
+_OpQuantConfig = qtyping.OpQuantizationConfig
+_AlgorithmName = quantizer.AlgorithmName
+_Numeric = Union[int, float]
+
+
+DEFAULT_ACTIVATION_QUANT_SETTING = _TensorQuantConfig(
+    num_bits=8,
+    symmetric=False,
+    granularity=qtyping.QuantGranularity.TENSORWISE,
+)
+DEFAULT_WEIGHT_QUANT_SETTING = _TensorQuantConfig(
+    num_bits=8,
+    symmetric=True,
+    granularity=qtyping.QuantGranularity.CHANNELWISE,
+)
+
+
+def get_static_activation_quant_setting(
+    num_bits: int, symmetric: bool
+) -> _TensorQuantConfig:
+  return _TensorQuantConfig(
+      num_bits=num_bits,
+      symmetric=symmetric,
+      granularity=qtyping.QuantGranularity.TENSORWISE,
+  )
+
+
+def get_static_op_quant_config(
+    activation_config: _TensorQuantConfig = DEFAULT_ACTIVATION_QUANT_SETTING,
+    weight_config: _TensorQuantConfig = DEFAULT_WEIGHT_QUANT_SETTING,
+) -> _OpQuantConfig:
+  return qtyping.OpQuantizationConfig(
+      activation_tensor_config=activation_config,
+      weight_tensor_config=weight_config,
+      compute_precision=_ComputePrecision.INTEGER,
+  )
+
 
 def get_path_to_datafile(path):
   """Get the path to the specified file in the data dependencies.
@@ -46,62 +89,152 @@ def get_path_to_datafile(path):
   return path
 
 
-def create_random_normal_dataset(
-    input_details: dict[str, Any],
-    num_samples: int,
-    random_seed: Union[int, np._typing.ArrayLike],
-) -> list[dict[str, Any]]:
-  """create random dataset following random distribution.
-
-  Args:
-    input_details: list of dict created by
-      tensorflow.lite.interpreter.get_input_details() for generating dataset
-    num_samples: number of input samples to be generated
-    random_seed: random seed to be used for function
-
-  Returns:
-    a list of inputs to the given interpreter, for a single interpreter we may
-    have multiple input tensors so each set of inputs is also represented as
-    list
-  """
-  rng = np.random.default_rng(random_seed)
-  dataset = []
-  for _ in range(num_samples):
-    input_data = {}
-    for arg_name, input_tensor in input_details.items():
-      new_data = rng.normal(size=input_tensor['shape']).astype(
-          input_tensor['dtype']
+class BaseOpTestCase(parameterized.TestCase):
+  """Base class for op-level tests."""
+
+  def quantize_and_validate(
+      self,
+      model_path: str,
+      algorithm_key: _AlgorithmName,
+      op_name: _OpName,
+      op_config: _OpQuantConfig,
+      num_validation_samples: int = 4,
+      num_calibration_samples: Optional[int] = None,
+      error_metric: str = 'mse',
+      min_max_range: Optional[tuple[_Numeric, _Numeric]] = None,
+  ) -> model_validator.ComparisonResult:
+    """Quantizes and validates the given model with the given configurations.
+
+    Args:
+      model_path: The path to the model to be quantized.
+      algorithm_key: The algorithm to be used for quantization.
+      op_name: The name of the operation to be quantized.
+      op_config: The configuration for the operation to be quantized.
+      num_validation_samples: The number of samples to use for validation.
+      num_calibration_samples: The number of samples to use for calibration. If
+        None then it will be set to num_validation_samples * 8.
+      error_metric: The error error_metric to use for validation.
+      min_max_range: The min and max of the input range.
+
+    Returns:
+      The comparison result of the validation.
+    """
+    quantizer_instance = quantizer.Quantizer(model_path)
+    quantizer_instance.update_quantization_recipe(
+        algorithm_key=algorithm_key,
+        regex='.*',
+        operation_name=op_name,
+        op_config=op_config,
+    )
+    if quantizer_instance.need_calibration:
+      if num_calibration_samples is None:
+        num_calibration_samples = num_validation_samples * 8
+      calibration_data = tfl_interpreter_utils.create_random_normal_input_data(
+          quantizer_instance.float_model,
+          num_samples=num_calibration_samples,
+          min_max_range=min_max_range,
+      )
+      calibration_result = quantizer_instance.calibrate(calibration_data)
+      quantization_result = quantizer_instance.quantize(calibration_result)
+    else:
+      quantization_result = quantizer_instance.quantize()
+    test_data = tfl_interpreter_utils.create_random_normal_input_data(
+        quantization_result.quantized_model,
+        num_samples=num_validation_samples,
+        min_max_range=min_max_range,
+    )
+    return quantizer_instance.validate(test_data, error_metric)
+
+  def assert_model_size_reduction_above_min_pct(
+      self,
+      validation_result: model_validator.ComparisonResult,
+      min_pct: float,
+  ):
+    """Checks the model size reduction (percentage) against the given expectation."""
+    _, reduction_pct = validation_result.get_model_size_reduction()
+    self.assertGreater(reduction_pct, min_pct)
+
+  def assert_weights_errors_below_tolerance(
+      self,
+      validation_result: model_validator.ComparisonResult,
+      weight_tolerance: float,
+  ):
+    """Checks the weight tensors' numerical behavior against the given tolerance."""
+    self.assertNotEmpty(validation_result.available_signature_keys())
+    for signature_key in validation_result.available_signature_keys():
+      signature_result = validation_result.get_signature_comparison_result(
+          signature_key
+      )
+      for result in signature_result.constant_tensors.values():
+        self.assertLess(result, weight_tolerance)
+
+  def assert_output_errors_below_tolerance(
+      self,
+      validation_result: model_validator.ComparisonResult,
+      output_tolerance: float,
+  ):
+    """Checks the output tensor numerical behavior against the given tolerance."""
+    self.assertNotEmpty(validation_result.available_signature_keys())
+    for signature_key in validation_result.available_signature_keys():
+      signature_result = validation_result.get_signature_comparison_result(
+          signature_key
+      )
+      for result in signature_result.output_tensors.values():
+        self.assertLess(result, output_tolerance)
+
+  def assert_quantization_accuracy_and_size(
+      self,
+      algorithm_key: _AlgorithmName,
+      model_path: str,
+      op_name: _OpName,
+      op_config: _OpQuantConfig,
+      expected_model_size_reduction: float,
+      weight_tolerance: float = 1e-4,
+      output_tolerance: float = 1e-4,
+      min_max_range: Optional[tuple[_Numeric, _Numeric]] = None,
+  ):
+    """Check if the quantization is successful and the result is valid."""
+    validation_result = self.quantize_and_validate(
+        model_path=model_path,
+        algorithm_key=algorithm_key,
+        op_name=op_name,
+        op_config=op_config,
+        min_max_range=min_max_range,
+    )
+    with self.subTest(name='ModelSizeReduction'):
+      self.assert_model_size_reduction_above_min_pct(
+          validation_result, expected_model_size_reduction
+      )
+    with self.subTest(name='WeightsErrors'):
+      self.assert_weights_errors_below_tolerance(
+          validation_result, weight_tolerance
+      )
+    with self.subTest(name='OutputErrors'):
+      self.assert_output_errors_below_tolerance(
+          validation_result, output_tolerance
       )
-      input_data[arg_name] = new_data
-    dataset.append(input_data)
-  return dataset
-
-
-def create_random_normal_input_data(
-    tflite_model: Union[str, bytes],
-    num_samples: int = 4,
-    random_seed: int = 666,
-) -> dict[str, list[dict[str, Any]]]:
-  """create random dataset following random distribution for signature runner.
-
-  Args:
-    tflite_model: TFLite model path or bytearray
-    num_samples: number of input samples to be generated
-    random_seed: random seed to be used for function
 
-  Returns:
-    a list of inputs to the given interpreter, for a single interpreter we may
-    have multiple signatures so each set of inputs is also represented as
-    list
-  """
-  tfl_interpreter = tfl_interpreter_utils.create_tfl_interpreter(tflite_model)
-  signature_defs = tfl_interpreter.get_signature_list()
-  signature_keys = list(signature_defs.keys())
-  test_data = {}
-  for signature_key in signature_keys:
-    signature_runner = tfl_interpreter.get_signature_runner(signature_key)
-    input_details = signature_runner.get_input_details()
-    test_data[signature_key] = create_random_normal_dataset(
-        input_details, num_samples, random_seed
+  def assert_quantization_accuracy(
+      self,
+      algorithm_key: _AlgorithmName,
+      model_path: str,
+      op_name: _OpName,
+      op_config: _OpQuantConfig,
+      num_validation_samples: int = 4,
+      num_calibration_samples: Optional[int] = None,
+      output_tolerance: float = 1e-4,
+      min_max_range: Optional[tuple[_Numeric, _Numeric]] = None,
+  ):
+    """Checks if the output errors after quantization are within the tolerance."""
+    validation_result = self.quantize_and_validate(
+        model_path=model_path,
+        algorithm_key=algorithm_key,
+        num_validation_samples=num_validation_samples,
+        num_calibration_samples=num_calibration_samples,
+        op_name=op_name,
+        op_config=op_config,
+        min_max_range=min_max_range,
+    )
+    self.assert_output_errors_below_tolerance(
+        validation_result, output_tolerance
     )
-  return test_data

ai_edge_quantizer/utils/tfl_flatbuffer_utils.py

@@ -20,48 +20,66 @@ from typing import Any, Optional, Union
 import immutabledict
 import numpy as np
 
+from ai_edge_litert.tools import flatbuffer_utils
 from ai_edge_quantizer import qtyping
-from ai_edge_litert import schema_py_generated  # pylint:disable=g-direct-tensorflow-import
-from tensorflow.lite.tools import flatbuffer_utils  # pylint: disable=g-direct-tensorflow-import
-from tensorflow.python.platform import gfile  # pylint: disable=g-direct-tensorflow-import
+from ai_edge_litert import schema_py_generated as schema  # pylint:disable=g-direct-tensorflow-import
+import os # tensorflow.python.platform.gfile  # pylint: disable=g-direct-tensorflow-import
 
 _TFLOpName = qtyping.TFLOperationName
 
 TFL_OP_NAME_TO_CODE = immutabledict.immutabledict({
-    _TFLOpName.FULLY_CONNECTED: (
-        schema_py_generated.BuiltinOperator.FULLY_CONNECTED
+    _TFLOpName.FULLY_CONNECTED: schema.BuiltinOperator.FULLY_CONNECTED,
+    _TFLOpName.BATCH_MATMUL: schema.BuiltinOperator.BATCH_MATMUL,
+    _TFLOpName.CONV_2D: schema.BuiltinOperator.CONV_2D,
+    _TFLOpName.DEPTHWISE_CONV_2D: schema.BuiltinOperator.DEPTHWISE_CONV_2D,
+    _TFLOpName.CONV_2D_TRANSPOSE: schema.BuiltinOperator.TRANSPOSE_CONV,
+    _TFLOpName.EMBEDDING_LOOKUP: schema.BuiltinOperator.EMBEDDING_LOOKUP,
+    _TFLOpName.SOFTMAX: schema.BuiltinOperator.SOFTMAX,
+    _TFLOpName.AVERAGE_POOL_2D: schema.BuiltinOperator.AVERAGE_POOL_2D,
+    _TFLOpName.RESHAPE: schema.BuiltinOperator.RESHAPE,
+    _TFLOpName.TANH: schema.BuiltinOperator.TANH,
+    _TFLOpName.TRANSPOSE: schema.BuiltinOperator.TRANSPOSE,
+    _TFLOpName.GELU: schema.BuiltinOperator.GELU,
+    _TFLOpName.ADD: schema.BuiltinOperator.ADD,
+    _TFLOpName.SUB: schema.BuiltinOperator.SUB,
+    _TFLOpName.MUL: schema.BuiltinOperator.MUL,
+    _TFLOpName.MEAN: schema.BuiltinOperator.MEAN,
+    _TFLOpName.RSQRT: schema.BuiltinOperator.RSQRT,
+    _TFLOpName.CONCATENATION: schema.BuiltinOperator.CONCATENATION,
+    _TFLOpName.STRIDED_SLICE: schema.BuiltinOperator.STRIDED_SLICE,
+    _TFLOpName.SPLIT: schema.BuiltinOperator.SPLIT,
+    _TFLOpName.LOGISTIC: schema.BuiltinOperator.LOGISTIC,
+    _TFLOpName.SLICE: schema.BuiltinOperator.SLICE,
+    _TFLOpName.SUM: schema.BuiltinOperator.SUM,
+    _TFLOpName.SELECT: schema.BuiltinOperator.SELECT,
+    _TFLOpName.SELECT_V2: schema.BuiltinOperator.SELECT_V2,
+    _TFLOpName.STABLEHLO_COMPOSITE: schema.BuiltinOperator.STABLEHLO_COMPOSITE,
+    _TFLOpName.DYNAMIC_UPDATE_SLICE: (
+        schema.BuiltinOperator.DYNAMIC_UPDATE_SLICE
     ),
-    _TFLOpName.BATCH_MATMUL: schema_py_generated.BuiltinOperator.BATCH_MATMUL,
-    _TFLOpName.CONV_2D: schema_py_generated.BuiltinOperator.CONV_2D,
-    _TFLOpName.DEPTHWISE_CONV_2D: (
-        schema_py_generated.BuiltinOperator.DEPTHWISE_CONV_2D
+    _TFLOpName.PAD: schema.BuiltinOperator.PAD,
+    _TFLOpName.SQUARED_DIFFERENCE: schema.BuiltinOperator.SQUARED_DIFFERENCE,
+    _TFLOpName.MAX_POOL_2D: schema.BuiltinOperator.MAX_POOL_2D,
+    _TFLOpName.RESIZE_BILINEAR: schema.BuiltinOperator.RESIZE_BILINEAR,
+    _TFLOpName.RESIZE_NEAREST_NEIGHBOR: (
+        schema.BuiltinOperator.RESIZE_NEAREST_NEIGHBOR
     ),
-    _TFLOpName.CONV_2D_TRANSPOSE: (
-        schema_py_generated.BuiltinOperator.TRANSPOSE_CONV
-    ),
-    _TFLOpName.EMBEDDING_LOOKUP: (
-        schema_py_generated.BuiltinOperator.EMBEDDING_LOOKUP
-    ),
-    _TFLOpName.SOFTMAX: schema_py_generated.BuiltinOperator.SOFTMAX,
-    _TFLOpName.AVERAGE_POOL_2D: (
-        schema_py_generated.BuiltinOperator.AVERAGE_POOL_2D
-    ),
-    _TFLOpName.RESHAPE: schema_py_generated.BuiltinOperator.RESHAPE,
-    _TFLOpName.TANH: schema_py_generated.BuiltinOperator.TANH,
-    _TFLOpName.TRANSPOSE: schema_py_generated.BuiltinOperator.TRANSPOSE,
-    _TFLOpName.GELU: schema_py_generated.BuiltinOperator.GELU,
-    _TFLOpName.ADD: schema_py_generated.BuiltinOperator.ADD,
-    _TFLOpName.SUB: schema_py_generated.BuiltinOperator.SUB,
-    _TFLOpName.MUL: schema_py_generated.BuiltinOperator.MUL,
-    _TFLOpName.MEAN: schema_py_generated.BuiltinOperator.MEAN,
-    _TFLOpName.RSQRT: schema_py_generated.BuiltinOperator.RSQRT,
-    _TFLOpName.CONCATENATION: schema_py_generated.BuiltinOperator.CONCATENATION,
-    _TFLOpName.STRIDED_SLICE: schema_py_generated.BuiltinOperator.STRIDED_SLICE,
-    _TFLOpName.SPLIT: schema_py_generated.BuiltinOperator.SPLIT,
-    _TFLOpName.LOGISTIC: schema_py_generated.BuiltinOperator.LOGISTIC,
-    _TFLOpName.SLICE: schema_py_generated.BuiltinOperator.SLICE,
-    _TFLOpName.SUM: schema_py_generated.BuiltinOperator.SUM,
-    _TFLOpName.SELECT_V2: schema_py_generated.BuiltinOperator.SELECT_V2,
+    _TFLOpName.GATHER_ND: schema.BuiltinOperator.GATHER_ND,
+    _TFLOpName.PACK: schema.BuiltinOperator.PACK,
+    _TFLOpName.UNPACK: schema.BuiltinOperator.UNPACK,
+    _TFLOpName.DIV: schema.BuiltinOperator.DIV,
+    _TFLOpName.BROADCAST_TO: schema.BuiltinOperator.BROADCAST_TO,
+    _TFLOpName.SQRT: schema.BuiltinOperator.SQRT,
+    _TFLOpName.GATHER: schema.BuiltinOperator.GATHER,
+    _TFLOpName.HARD_SWISH: schema.BuiltinOperator.HARD_SWISH,
+    _TFLOpName.MAXIMUM: schema.BuiltinOperator.MAXIMUM,
+    _TFLOpName.PADV2: schema.BuiltinOperator.PADV2,
+    _TFLOpName.REDUCE_MIN: schema.BuiltinOperator.REDUCE_MIN,
+    _TFLOpName.EQUAL: schema.BuiltinOperator.EQUAL,
+    _TFLOpName.NOT_EQUAL: schema.BuiltinOperator.NOT_EQUAL,
+    _TFLOpName.MIRROR_PAD: schema.BuiltinOperator.MIRROR_PAD,
+    _TFLOpName.SPACE_TO_DEPTH: schema.BuiltinOperator.SPACE_TO_DEPTH,
+    _TFLOpName.RELU: schema.BuiltinOperator.RELU,
 })
 
 TFL_OP_CODE_TO_NAME = immutabledict.immutabledict(
@@ -78,6 +96,11 @@ TFL_OP_TO_WEIGHT_QUANTIZED_DIM = immutabledict.immutabledict({
     _TFLOpName.CONV_2D_TRANSPOSE: 0,
 })
 
+TFL_OP_TO_BLOCKWISE_WEIGHT_QUANTIZED_DIM = immutabledict.immutabledict({
+    _TFLOpName.FULLY_CONNECTED: 1,
+    _TFLOpName.EMBEDDING_LOOKUP: 1,
+})
+
 NUM_TFL_DATATYPES = 18
 TENSOR_CODE_TO_TYPE = {}
 for dtype_code in range(NUM_TFL_DATATYPES):
@@ -87,7 +110,7 @@ TENSOR_TYPE_TO_CODE = immutabledict.immutabledict(
     (reversed(item) for item in TENSOR_CODE_TO_TYPE.items())
 )
 
-# Expose functions in tensorflow.lite.tools.flatbuffer_utils
+# Expose functions in litert.python.tools.flatbuffer_utils
 write_model = flatbuffer_utils.write_model
 
 
@@ -122,7 +145,7 @@ def get_model_content(tflite_path: str) -> bytes:
   Returns:
     The model bytes.
   """
-  with gfile.Open(tflite_path, "rb") as tflite_file:
+  with open(tflite_path, "rb") as tflite_file:
     return tflite_file.read()
 
 
@@ -135,7 +158,7 @@ def get_model_buffer(tflite_path: str) -> bytearray:
   Returns:
     model_buffer: the model buffer.
   """
-  with gfile.Open(tflite_path, "rb") as tflite_file:
+  with open(tflite_path, "rb") as tflite_file:
     return bytearray(tflite_file.read())
 
 
@@ -188,25 +211,18 @@ def parse_fc_bmm_conv_tensors(
   return input_tensor, weight_tensor, bias_tensor, output_tensor
 
 
-# flatbuffer_model has Any type since tensorflow.lite.tools.flatbuffer_utils
+# flatbuffer_model has Any type since litert.python.tools.flatbuffer_utils
 # is not type annotated.
 def buffer_to_tensors(flatbuffer_model: Any) -> dict[int, list[Any]]:
-  """Get the buffer to tensor map for a tflite model.
-
-  Args:
-    flatbuffer_model: the flatbuffer_model.
-
-  Returns:
-    buffer_to_tensor_map: key as buffer index, value as list of tensors share
-    the buffer
-  """
+  """Returns a map from buffer id to tensors that use it."""
   buffer_to_tensor_map = {}
   for subgraph in flatbuffer_model.subgraphs:
     for op in subgraph.operators:
       for tensor in parse_op_tensors(op, subgraph.tensors):
         if tensor.buffer not in buffer_to_tensor_map:
           buffer_to_tensor_map[tensor.buffer] = []
-        buffer_to_tensor_map[tensor.buffer].append(tensor)
+        if tensor not in buffer_to_tensor_map[tensor.buffer]:
+          buffer_to_tensor_map[tensor.buffer].append(tensor)
   return buffer_to_tensor_map
 
 
@@ -239,7 +255,8 @@ def get_tensor_data(tensor: Any, buffers: list[Any]) -> Optional[np.ndarray]:
   data = np.frombuffer(
       buffer_data, dtype=TENSOR_CODE_TO_TYPE[tensor.type].lower()
   )
-  data = np.reshape(data, tensor.shape)
+  if tensor.shape is not None:
+    data = np.reshape(data, tensor.shape)
   return data
 
 
@@ -315,3 +332,32 @@ def get_subgraph_input_output_operators(
       op_key=qtyping.TFLOperationName.OUTPUT,
   )
   return [input_op, output_op]
+
+
+def get_op_side_effect_subgraphs(
+    op: Union[schema.Operator, schema.OperatorT],
+) -> list[int]:
+  """Get indices of any subgraphs invoked as a side effect of the operator.
+
+  Args:
+    op: The operator object.
+
+  Returns:
+    A list of subgraph indices invoked by the operator. Empty if the operator
+    does not invoke any subgraphs.
+  """
+  if opts := flatbuffer_utils.get_options_as(
+      op, schema.StableHLOCompositeOptionsT
+  ):
+    return [opts.decompositionSubgraphIndex]
+  # Can add other nested ops here (control flow ops, etc).
+  return []
+
+
+def get_op_name_by_index(
+    flatbuffer_model: Any, subgraph_id: int, op_index: int
+) -> str:
+  """Get the op name from the flatbuffer model."""
+  op = flatbuffer_model.subgraphs[subgraph_id].operators[op_index]
+  builtin_code = flatbuffer_model.operatorCodes[op.opcodeIndex].builtinCode
+  return TFL_OP_CODE_TO_NAME[builtin_code]

ai_edge_quantizer/utils/tfl_flatbuffer_utils_test.py

@@ -105,6 +105,26 @@ class FlatbufferUtilsTest(googletest.TestCase):
     conv2d_filter_tensor = tensors[0]
     self.assertEqual(tuple(conv2d_filter_tensor.shape), (8, 3, 3, 1))
 
+  def test_buffer_to_tensors_has_unique_values(self):
+    test_model_path = os.path.join(
+        TEST_DATA_PREFIX_PATH,
+        "constant_tensor_and_buffer_only_sharing_weight_fcs.tflite",
+    )
+    test_model = tfl_flatbuffer_utils.read_model(test_model_path)
+    buffer_to_tensor_map = tfl_flatbuffer_utils.buffer_to_tensors(test_model)
+    self.assertLen(buffer_to_tensor_map, 7)
+    # The following buffer is shared by two tensors, each shared by two FC ops.
+    # This is where before we had multiple enrties for the same tensor.
+    self.assertLen(buffer_to_tensor_map[2], 2)
+    got_tensor_names = [
+        tfl_flatbuffer_utils.get_tensor_name(tensor)
+        for tensor in buffer_to_tensor_map[2]
+    ]
+    self.assertEqual(
+        got_tensor_names,
+        ["arith.constant", "arith.constant1"],
+    )
+
   def test_get_tensor_name(self):
     subgraph0 = self._test_model.subgraphs[0]
     subgraph_tensors = subgraph0.tensors