ai-edge-quantizer-nightly 0.0.1.dev20250115__py3-none-any.whl

ai_edge_quantizer/utils/tfl_interpreter_utils.py

@@ -0,0 +1,312 @@
+# 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.
+# ==============================================================================
+
+"""Util functions for TFL interpreter."""
+
+from typing import Any, Optional, Union
+
+import numpy as np
+
+from ai_edge_quantizer import qtyping
+from ai_edge_quantizer.algorithms.uniform_quantize import uniform_quantize_tensor
+from ai_edge_litert import interpreter as tfl  # pylint: disable=g-direct-tensorflow-import
+from tensorflow.python.platform import gfile  # pylint: disable=g-direct-tensorflow-import
+
+DEFAULT_SIGNATURE_KEY = "serving_default"
+
+
+def create_tfl_interpreter(
+    tflite_model: Union[str, bytes],
+    allocate_tensors: bool = True,
+    use_xnnpack: bool = True,
+    num_threads: int = 16,
+) -> tfl.Interpreter:
+  """Creates a TFLite interpreter from a model file.
+
+  Args:
+    tflite_model: Model file path or bytes.
+    allocate_tensors: Whether to allocate tensors.
+    use_xnnpack: Whether to use the XNNPACK delegate for the interpreter.
+    num_threads: The number of threads to use for the interpreter.
+
+  Returns:
+    A TFLite interpreter.
+  """
+  if isinstance(tflite_model, str):
+    with gfile.GFile(tflite_model, "rb") as f:
+      tflite_model = f.read()
+
+  if use_xnnpack:
+    op_resolver = tfl.OpResolverType.BUILTIN
+  else:
+    op_resolver = tfl.OpResolverType.BUILTIN_WITHOUT_DEFAULT_DELEGATES
+  tflite_interpreter = tfl.Interpreter(
+      model_content=bytes(tflite_model),
+      num_threads=num_threads,
+      experimental_op_resolver_type=op_resolver,
+      experimental_preserve_all_tensors=True,
+  )
+  if allocate_tensors:
+    tflite_interpreter.allocate_tensors()
+  return tflite_interpreter
+
+
+def is_tensor_quantized(tensor_detail: dict[str, Any]) -> bool:
+  """Checks if a tensor is quantized.
+
+  Args:
+    tensor_detail: A dictionary of tensor details.
+
+  Returns:
+    True if the tensor is quantized.
+  """
+  quant_params = tensor_detail["quantization_parameters"]
+  return bool(len(quant_params["scales"]))
+
+
+def invoke_interpreter_signature(
+    tflite_interpreter: tfl.Interpreter,
+    signature_input_data: dict[str, Any],
+    signature_key: Optional[str] = None,
+    quantize_input: bool = True,
+) -> dict[str, np.ndarray]:
+  """Invokes the TFLite interpreter through signature runner.
+
+  Args:
+    tflite_interpreter: A TFLite interpreter.
+    signature_input_data: The input data for the signature.
+    signature_key: The signature key.
+    quantize_input: Whether to quantize the input data.
+
+  Returns:
+    The output data of the signature.
+  """
+  # Make a copy to avoid in-place modification.
+  signature_input = signature_input_data.copy()
+  signature_runner = tflite_interpreter.get_signature_runner(signature_key)
+  for input_name, input_detail in signature_runner.get_input_details().items():
+    if is_tensor_quantized(input_detail) and quantize_input:
+      input_data = signature_input[input_name]
+      quant_params = qtyping.UniformQuantParams.from_tfl_tensor_details(
+          input_detail
+      )
+      signature_input[input_name] = uniform_quantize_tensor.uniform_quantize(
+          input_data, quant_params
+      )
+  return signature_runner(**signature_input)
+
+
+def invoke_interpreter_once(
+    tflite_interpreter: tfl.Interpreter,
+    input_data_list: list[Any],
+    quantize_input: bool = True,
+):
+  """Invokes the TFLite interpreter once.
+
+  Args:
+    tflite_interpreter: A TFLite interpreter.
+    input_data_list: A list of input data.
+    quantize_input: Whether to quantize the input data.
+  """
+  if len(input_data_list) != len(tflite_interpreter.get_input_details()):
+    raise ValueError(
+        "Input data must be a list with each element match the input sequence"
+        " defined in .tflite. If the model has only one input, wrap it with a"
+        " list (e.g., [input_data])"
+    )
+  for i, input_data in enumerate(input_data_list):
+    input_details = tflite_interpreter.get_input_details()[i]
+    if is_tensor_quantized(input_details) and quantize_input:
+      quant_params = qtyping.UniformQuantParams.from_tfl_tensor_details(
+          input_details
+      )
+      input_data = uniform_quantize_tensor.uniform_quantize(
+          input_data, quant_params
+      )
+    tflite_interpreter.set_tensor(input_details["index"], input_data)
+  tflite_interpreter.invoke()
+
+
+def get_tensor_data(
+    tflite_interpreter: Any,
+    tensor_detail: dict[str, Any],
+    subgraph_index: int = 0,
+    dequantize: bool = True,
+) -> np.ndarray:
+  """Gets the tensor data from a TFLite interpreter.
+
+  Args:
+    tflite_interpreter: A TFLite interpreter.
+    tensor_detail: A dictionary of tensor details.
+    subgraph_index: The index of the subgraph that the tensor belongs to.
+    dequantize: Whether to dequantize the quantized tensor data.
+
+  Returns:
+    The tensor data.
+  """
+  tensor_data = tflite_interpreter.get_tensor(
+      tensor_detail["index"], subgraph_index
+  )
+  if is_tensor_quantized(tensor_detail) and dequantize:
+    quant_params = qtyping.UniformQuantParams.from_tfl_tensor_details(
+        tensor_detail
+    )
+    tensor_data = uniform_quantize_tensor.uniform_dequantize(
+        tensor_data,
+        quant_params,
+    )
+  return tensor_data
+
+
+def get_tensor_name_to_content_map(
+    tflite_interpreter: Any, subgraph_index: int = 0, dequantize: bool = False
+) -> dict[str, Any]:
+  """Gets internal tensors from a TFLite interpreter for a given subgraph.
+
+  Note the data will be copied to the returned dictionary, increasing the
+  memory usage.
+
+  Args:
+    tflite_interpreter: A TFLite interpreter.
+    subgraph_index: The index of the subgraph that the tensor belongs to.
+    dequantize: Whether to dequantize the tensor data.
+
+  Returns:
+    A dictionary of internal tensors.
+  """
+  tensors = {}
+  for tensor_detail in tflite_interpreter.get_tensor_details(subgraph_index):
+    # Don't return temporary, unnamed tensors
+    if not tensor_detail["name"]:
+      continue
+    tensors[tensor_detail["name"]] = get_tensor_data(
+        tflite_interpreter, tensor_detail, subgraph_index, dequantize
+    )
+  return tensors
+
+
+def get_tensor_name_to_details_map(
+    tflite_interpreter: Any, subgraph_index: int = 0
+) -> dict[str, Any]:
+  """Gets internal tensors from a TFLite interpreter for a given subgraph.
+
+  Args:
+    tflite_interpreter: A TFLite interpreter.
+    subgraph_index: The index of the subgraph that the tensor belongs to.
+
+  Returns:
+    A dictionary of internal tensors.
+  """
+  tensor_name_to_detail = {}
+  for tensor_detail in tflite_interpreter.get_tensor_details(subgraph_index):
+    # Don't return temporary, unnamed tensors
+    if not tensor_detail["name"]:
+      continue
+    tensor_name_to_detail[tensor_detail["name"]] = tensor_detail
+  return tensor_name_to_detail
+
+
+def get_input_tensor_names(
+    tflite_model: Union[str, bytes], signature_name: Optional[str] = None
+) -> list[str]:
+  """Gets input tensor names from a TFLite model for a signature.
+
+  Args:
+    tflite_model: Model file path or bytes.
+    signature_name: The signature name that the input tensors belong to.
+
+  Returns:
+    A list of input tensor names.
+  """
+
+  tfl_interpreter = create_tfl_interpreter(tflite_model, allocate_tensors=False)
+  signature_runner = tfl_interpreter.get_signature_runner(signature_name)
+  input_tensor_names = []
+  for _, input_detail in signature_runner.get_input_details().items():
+    input_tensor_names.append(input_detail["name"])
+  return input_tensor_names
+
+
+def get_output_tensor_names(
+    tflite_model: Union[str, bytes], signature_name: Optional[str] = None
+) -> list[str]:
+  """Gets output tensor names from a TFLite model for a signature.
+
+  Args:
+    tflite_model: Model file path or bytes.
+    signature_name: The signature name that the output tensors belong to.
+
+  Returns:
+    A list of output tensor names.
+  """
+  tfl_interpreter = create_tfl_interpreter(tflite_model, allocate_tensors=False)
+  signature_runner = tfl_interpreter.get_signature_runner(signature_name)
+  output_tensor_names = []
+  for _, output_detail in signature_runner.get_output_details().items():
+    output_tensor_names.append(output_detail["name"])
+  return output_tensor_names
+
+
+def get_constant_tensor_names(
+    tflite_model: Union[str, bytes],
+    subgraph_index: int = 0,
+    min_constant_size: int = 1,
+) -> list[str]:
+  """Gets constant tensor names from a TFLite model for a subgraph.
+
+  Note that this function acts on subgraph level, not signature level. This is
+  because it is non-trivial to track constant tensors for a signature without
+  running it.
+
+  Args:
+    tflite_model: Model file path or bytes.
+    subgraph_index: The index of the subgraph that the tensor belongs to.
+    min_constant_size: The minimum size of a constant tensor.
+
+  Returns:
+    A list of names for constant tensor that bigger than min_constant_size and a
+    list of names for constant tensor that smaller than min_constant_size.
+  """
+  tfl_interpreter = create_tfl_interpreter(tflite_model, allocate_tensors=False)
+  const_tensor_names = []
+  for tensor_detail in tfl_interpreter.get_tensor_details(subgraph_index):
+    if tensor_detail["dtype"] == np.object_:
+      continue
+    try:
+      tensor_data = get_tensor_data(
+          tfl_interpreter, tensor_detail, subgraph_index
+      )
+      if tensor_data.size >= min_constant_size:
+        const_tensor_names.append(tensor_detail["name"])
+    except ValueError:
+      continue
+  return const_tensor_names
+
+
+def get_signature_main_subgraph_index(
+    tflite_interpreter: tfl.Interpreter,
+    signature_key: Optional[str] = None,
+) -> int:
+  """Gets the main subgraph index of a signature.
+
+  Args:
+    tflite_interpreter: A TFLite interpreter.
+    signature_key: The signature key.
+
+  Returns:
+    The main subgraph index of the signature.
+  """
+  signature_runner = tflite_interpreter.get_signature_runner(signature_key)
+  return signature_runner._subgraph_index  # pylint:disable=protected-access

ai_edge_quantizer/utils/tfl_interpreter_utils_test.py

@@ -0,0 +1,332 @@
+# 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.
+# ==============================================================================
+
+import os
+import numpy as np
+from tensorflow.python.platform import googletest
+from ai_edge_quantizer.utils import test_utils
+from ai_edge_quantizer.utils import tfl_interpreter_utils
+
+
+TEST_DATA_PREFIX_PATH = test_utils.get_path_to_datafile("../tests/models")
+
+
+class TflUtilsSingleSignatureModelTest(googletest.TestCase):
+
+  def setUp(self):
+    super().setUp()
+    np.random.seed(0)
+    self._test_model_path = os.path.join(
+        TEST_DATA_PREFIX_PATH, "conv_fc_mnist.tflite"
+    )
+    self._input_data = np.random.rand(1, 28, 28, 1).astype(np.float32)
+
+  def test_create_tfl_interpreter(self):
+    tfl_interpreter = tfl_interpreter_utils.create_tfl_interpreter(
+        self._test_model_path
+    )
+    self.assertIsNotNone(tfl_interpreter)
+
+  def test_invoke_interpreter_once(self):
+    tfl_interpreter = tfl_interpreter_utils.create_tfl_interpreter(
+        self._test_model_path
+    )
+    tfl_interpreter_utils.invoke_interpreter_once(
+        tfl_interpreter, [self._input_data]
+    )
+    output_details = tfl_interpreter.get_output_details()[0]
+    output_data = tfl_interpreter.get_tensor(output_details["index"])
+    self.assertIsNotNone(output_data)
+    self.assertEqual(tuple(output_data.shape), (1, 10))
+    self.assertAlmostEqual(output_data[0][0], 0.0031010755)
+
+  def test_get_tensor_data(self):
+    tfl_interpreter = tfl_interpreter_utils.create_tfl_interpreter(
+        self._test_model_path
+    )
+    tfl_interpreter_utils.invoke_interpreter_once(
+        tfl_interpreter, [self._input_data]
+    )
+    output_details = tfl_interpreter.get_output_details()[0]
+    output_data = tfl_interpreter_utils.get_tensor_data(
+        tfl_interpreter, output_details
+    )
+    self.assertEqual(tuple(output_data.shape), (1, 10))
+
+  def test_get_tensor_name_to_content_map(self):
+    tfl_interpreter = tfl_interpreter_utils.create_tfl_interpreter(
+        self._test_model_path
+    )
+    tfl_interpreter_utils.invoke_interpreter_once(
+        tfl_interpreter, [self._input_data]
+    )
+
+    tensor_name_to_content_map = (
+        tfl_interpreter_utils.get_tensor_name_to_content_map(tfl_interpreter)
+    )
+    input_content = tensor_name_to_content_map["serving_default_conv2d_input:0"]
+    self.assertSequenceAlmostEqual(
+        self._input_data.flatten(), input_content.flatten()
+    )
+    weight_content = tensor_name_to_content_map["sequential/conv2d/Conv2D"]
+    self.assertEqual(tuple(weight_content.shape), (8, 3, 3, 1))
+
+    self.assertIn(
+        "sequential/average_pooling2d/AvgPool", tensor_name_to_content_map
+    )
+    average_pool_res = tensor_name_to_content_map[
+        "sequential/average_pooling2d/AvgPool"
+    ]
+    self.assertEqual(tuple(average_pool_res.shape), (1, 14, 14, 8))
+
+  def test_is_tensor_quantized(self):
+    tfl_interpreter = tfl_interpreter_utils.create_tfl_interpreter(
+        self._test_model_path
+    )
+    input_details = tfl_interpreter.get_input_details()[0]
+    self.assertFalse(tfl_interpreter_utils.is_tensor_quantized(input_details))
+
+  def test_get_input_tensor_names(self):
+    input_tensor_names = tfl_interpreter_utils.get_input_tensor_names(
+        self._test_model_path
+    )
+    self.assertEqual(
+        input_tensor_names,
+        ["serving_default_conv2d_input:0"],
+    )
+
+  def test_get_output_tensor_names(self):
+    output_tensor_names = tfl_interpreter_utils.get_output_tensor_names(
+        self._test_model_path
+    )
+    self.assertEqual(
+        output_tensor_names,
+        ["StatefulPartitionedCall:0"],
+    )
+
+  def test_get_constant_tensor_names(self):
+    const_tensor_names = tfl_interpreter_utils.get_constant_tensor_names(
+        self._test_model_path
+    )
+    self.assertEqual(
+        set(const_tensor_names),
+        set([
+            "sequential/conv2d/Conv2D",
+            "sequential/conv2d/Relu;sequential/conv2d/BiasAdd;sequential/conv2d/Conv2D;sequential/conv2d/BiasAdd/ReadVariableOp",
+            "arith.constant",
+            "arith.constant1",
+            "arith.constant2",
+            "arith.constant3",
+        ]),
+    )
+
+
+class TflUtilsQuantizedModelTest(googletest.TestCase):
+
+  def setUp(self):
+    super().setUp()
+    np.random.seed(0)
+    self._test_model_path = os.path.join(
+        TEST_DATA_PREFIX_PATH, "conv_fc_mnist_srq_a8w8.tflite"
+    )
+    self._signature_input_data = {
+        "conv2d_input": np.random.rand(1, 28, 28, 1).astype(np.float32)
+    }
+
+  def test_is_tensor_quantized(self):
+    tfl_interpreter = tfl_interpreter_utils.create_tfl_interpreter(
+        self._test_model_path
+    )
+    input_details = tfl_interpreter.get_input_details()[0]
+    self.assertTrue(tfl_interpreter_utils.is_tensor_quantized(input_details))
+
+  def test_invoke_interpreter_signature(self):
+    tfl_interpreter = tfl_interpreter_utils.create_tfl_interpreter(
+        self._test_model_path
+    )
+    signature_output = tfl_interpreter_utils.invoke_interpreter_signature(
+        tfl_interpreter, self._signature_input_data
+    )
+    print(signature_output)
+    self.assertEqual(tuple(signature_output["dense_1"].shape), (1, 10))
+
+    # Assert the input data is not modified in-place b/353340272.
+    self.assertEqual(
+        self._signature_input_data["conv2d_input"].dtype, np.float32
+    )
+
+
+class TflUtilsMultiSignatureModelTest(googletest.TestCase):
+
+  def setUp(self):
+    super().setUp()
+    np.random.seed(0)
+    self._test_model_path = os.path.join(
+        TEST_DATA_PREFIX_PATH, "two_signatures.tflite"
+    )
+    self._signature_input_data = {"x": np.array([2.0]).astype(np.float32)}
+
+  def test_create_tfl_interpreter(self):
+    tfl_interpreter = tfl_interpreter_utils.create_tfl_interpreter(
+        self._test_model_path
+    )
+    self.assertIsNotNone(tfl_interpreter)
+
+  def test_get_input_tensor_names(self):
+    signature_name = "add"
+    input_tensor_names = tfl_interpreter_utils.get_input_tensor_names(
+        self._test_model_path, signature_name
+    )
+    self.assertEqual(
+        input_tensor_names,
+        ["add_x:0"],
+    )
+
+    signature_name = "multiply"
+    input_tensor_names = tfl_interpreter_utils.get_input_tensor_names(
+        self._test_model_path, signature_name
+    )
+    self.assertEqual(
+        input_tensor_names,
+        ["multiply_x:0"],
+    )
+
+  def test_get_output_tensor_names(self):
+    signature_name = "add"
+    input_tensor_names = tfl_interpreter_utils.get_output_tensor_names(
+        self._test_model_path, signature_name
+    )
+    self.assertEqual(
+        input_tensor_names,
+        ["PartitionedCall:0"],
+    )
+
+    signature_name = "multiply"
+    input_tensor_names = tfl_interpreter_utils.get_output_tensor_names(
+        self._test_model_path, signature_name
+    )
+    self.assertEqual(
+        input_tensor_names,
+        ["PartitionedCall_1:0"],
+    )
+
+  def test_get_constant_tensor_names(self):
+    subgraph0_const_tensor_names = (
+        tfl_interpreter_utils.get_constant_tensor_names(
+            self._test_model_path, 0
+        )
+    )
+    self.assertEqual(subgraph0_const_tensor_names, ["Add/y"])
+
+    subgraph1_const_tensor_names = (
+        tfl_interpreter_utils.get_constant_tensor_names(
+            self._test_model_path, 1
+        )
+    )
+    self.assertEqual(subgraph1_const_tensor_names, ["Mul/y"])
+
+  def test_get_signature_main_subgraph_index(self):
+    tfl_interpreter = tfl_interpreter_utils.create_tfl_interpreter(
+        self._test_model_path
+    )
+    add_subgraph_index = (
+        tfl_interpreter_utils.get_signature_main_subgraph_index(
+            tfl_interpreter, "add"
+        )
+    )
+    self.assertEqual(add_subgraph_index, 0)
+    multiply_subgraph_index = (
+        tfl_interpreter_utils.get_signature_main_subgraph_index(
+            tfl_interpreter, "multiply"
+        )
+    )
+    self.assertEqual(multiply_subgraph_index, 1)
+
+  def test_get_tensor_data(self):
+    tfl_interpreter = tfl_interpreter_utils.create_tfl_interpreter(
+        self._test_model_path
+    )
+    # Invoke the ADD signature.
+    tfl_interpreter_utils.invoke_interpreter_signature(
+        tfl_interpreter, self._signature_input_data, "add"
+    )
+    output_details = {"index": 2, "quantization_parameters": {"scales": []}}
+    output_data = tfl_interpreter_utils.get_tensor_data(
+        tfl_interpreter, output_details, subgraph_index=0
+    )  # The ADD signature is in the first subgraph.
+    self.assertEqual(output_data, [12.0])  # 10 + 2
+
+    # Invoke the MULTIPLY signature.
+    tfl_interpreter_utils.invoke_interpreter_signature(
+        tfl_interpreter, self._signature_input_data, "multiply"
+    )
+    output_data = tfl_interpreter_utils.get_tensor_data(
+        tfl_interpreter, output_details, subgraph_index=1
+    )  # The Multiply signature is in the second subgraph.
+    self.assertEqual(output_data, [20.0])  # 10 * 2
+
+  def test_get_tensor_name_to_content_map(self):
+    tfl_interpreter = tfl_interpreter_utils.create_tfl_interpreter(
+        self._test_model_path
+    )
+    # Invoke all signatures.
+    tfl_interpreter_utils.invoke_interpreter_signature(
+        tfl_interpreter, self._signature_input_data, "multiply"
+    )
+    tfl_interpreter_utils.invoke_interpreter_signature(
+        tfl_interpreter, self._signature_input_data, "add"
+    )
+
+    # Test tensors belonging to the ADD signature.
+    add_subgraph_index = (
+        tfl_interpreter_utils.get_signature_main_subgraph_index(
+            tfl_interpreter, "add"
+        )
+    )
+    add_tensor_content = tfl_interpreter_utils.get_tensor_name_to_content_map(
+        tfl_interpreter, add_subgraph_index
+    )
+
+    add_input_content = add_tensor_content["add_x:0"]
+    self.assertSequenceAlmostEqual(
+        self._signature_input_data["x"].flatten(), add_input_content.flatten()
+    )
+    weight_content = add_tensor_content["Add/y"]
+    self.assertEqual(weight_content, 10)
+    add_output_content = add_tensor_content["PartitionedCall:0"]
+    self.assertEqual(add_output_content, [12.0])
+
+    # Test tensors belonging to the MULTIPLY signature.
+    multiply_subgraph_index = (
+        tfl_interpreter_utils.get_signature_main_subgraph_index(
+            tfl_interpreter, "multiply"
+        )
+    )
+    mul_tensor_content = tfl_interpreter_utils.get_tensor_name_to_content_map(
+        tfl_interpreter, multiply_subgraph_index
+    )
+    multiply_input_content = mul_tensor_content["multiply_x:0"]
+    self.assertSequenceAlmostEqual(
+        self._signature_input_data["x"].flatten(),
+        multiply_input_content.flatten(),
+    )
+    weight_content = mul_tensor_content["Mul/y"]
+    self.assertEqual(weight_content, 10)
+    multiply_output_content = mul_tensor_content["PartitionedCall_1:0"]
+    self.assertEqual(multiply_output_content, [20.0])
+
+
+if __name__ == "__main__":
+  googletest.main()