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

ai_edge_quantizer/transformations/transformation_utils.py

@@ -0,0 +1,132 @@
+# 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.
+# ==============================================================================
+
+"""Utility functions for graph transformations."""
+
+import dataclasses
+from typing import Union
+
+import numpy as np
+
+from ai_edge_quantizer import qtyping
+from ai_edge_litert import schema_py_generated  # pylint: disable=g-direct-tensorflow-import
+
+
+@dataclasses.dataclass
+class TransformationInput:
+  """Standard input for a graph transformation.
+
+  Attributes:
+    tensor_id: the tensor index to insert dequant op after
+    op_codes: list of operatorCode in the model, if dequantize op doesn't exist,
+      we need to insert the op code into the list
+    buffers: list of buffer in the original TFlite model for buffer quantization
+    subgraph: flatbuffer subgraph object which the tensor resides.
+    producer: op id for the producer of the tensor.
+    consumers: op ids for consumers of the new dequant op.
+    quant_params: quantization parameters to be applied on the orignal tensor
+  """
+
+  tensor_id: int
+  op_codes: list[schema_py_generated.OperatorCodeT]
+  buffers: list[schema_py_generated.BufferT]
+  subgraph: schema_py_generated.SubGraphT
+  producer: int
+  consumers: list[int]
+  quant_params: Union[qtyping.UniformQuantParams, qtyping.NonLinearQuantParams]
+
+
+def add_op_code(
+    op_code: schema_py_generated.OperatorCodeT,
+    model_op_codes: list[schema_py_generated.OperatorCodeT],
+) -> int:
+  """Add an op code into a model if it's not present.
+
+  Args:
+    op_code: The op code to be added.
+    model_op_codes: The op codes of the model.
+
+  Returns:
+    The index of the op code in the model.
+  """
+  for i, model_op_code in enumerate(model_op_codes):
+    if model_op_code.builtinCode == op_code:
+      return i
+  model_op_codes.append(schema_py_generated.OperatorCodeT())
+  model_op_codes[-1].builtinCode = op_code
+  return len(model_op_codes) - 1
+
+
+def add_new_constant_tensor(
+    tensor_name: str,
+    data: np.ndarray,
+    tensor_type: schema_py_generated.TensorType,
+    subgraph: schema_py_generated.SubGraphT,
+    buffers: list[schema_py_generated.BufferT],
+) -> int:
+  """Add a new constant tensor to the model.
+
+  Args:
+    tensor_name: The name of the new tensor.
+    data: The data of the new tensor.
+    tensor_type: The type of the new tensor.
+    subgraph: The subgraph where the new tensor is added.
+    buffers: The buffers of the model.
+
+  Returns:
+    The index of the new tensor in the subgraph.
+  """
+  tensor_buffer = schema_py_generated.BufferT()
+  tensor_buffer.data = np.frombuffer(data.tobytes(), dtype=np.uint8).flatten()
+  tensor_buffer.offset = 0
+  tensor_buffer.size = 0
+  tensor_buffer_id = len(buffers)
+  buffers.append(tensor_buffer)
+
+  new_tensor = schema_py_generated.TensorT()
+  new_tensor.shape = data.shape
+  new_tensor.buffer = tensor_buffer_id
+  new_tensor.type = tensor_type
+  new_tensor.name = tensor_name
+  new_tensor_id = len(subgraph.tensors)
+  subgraph.tensors.append(new_tensor)
+  return new_tensor_id
+
+
+def add_new_activation_tensor(
+    tensor_name: str,
+    shape: list[int],
+    tensor_type: schema_py_generated.TensorType,
+    subgraph: schema_py_generated.SubGraphT,
+) -> int:
+  """Add a new activation tensor to the model.
+
+  Args:
+    tensor_name: The name of the new tensor.
+    shape: The shape of the new tensor.
+    tensor_type: The type of the new tensor.
+    subgraph: The subgraph where the new tensor is added.
+
+  Returns:
+    The index of the new tensor in the subgraph.
+  """
+  new_tensor = schema_py_generated.TensorT()
+  new_tensor.shape = shape
+  new_tensor.type = tensor_type
+  new_tensor.name = tensor_name
+  new_tensor.buffer = 0
+  new_tensor_id = len(subgraph.tensors)
+  subgraph.tensors.append(new_tensor)
+  return new_tensor_id

ai_edge_quantizer/transformations/transformation_utils_test.py

@@ -0,0 +1,162 @@
+# 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.
+# ==============================================================================
+
+"""Tests for transformation_utils."""
+
+import os
+import numpy as np
+from tensorflow.python.platform import googletest
+from absl.testing import parameterized
+from ai_edge_quantizer.transformations import transformation_utils
+from ai_edge_quantizer.utils import test_utils
+from ai_edge_quantizer.utils import tfl_flatbuffer_utils
+from ai_edge_litert import schema_py_generated  # pylint: disable=g-direct-tensorflow-import
+
+TEST_DATA_PREFIX_PATH = test_utils.get_path_to_datafile("../tests/models")
+
+
+class TransformationUtilsTest(parameterized.TestCase):
+
+  def setUp(self):
+    super().setUp()
+    self.model_path = os.path.join(
+        TEST_DATA_PREFIX_PATH, "single_fc_bias.tflite"
+    )
+    self.model = tfl_flatbuffer_utils.read_model(self.model_path)
+
+  @parameterized.named_parameters(
+      dict(
+          testcase_name="add_new_op_code",
+          op_code=schema_py_generated.BuiltinOperator.LOGISTIC,
+          expected=1,
+      ),
+      dict(
+          testcase_name="add_existing_op_code",
+          op_code=schema_py_generated.BuiltinOperator.FULLY_CONNECTED,
+          expected=0,
+      ),
+  )
+  def test_add_op_code(self, op_code, expected):
+    """Tests if the op code is added to the model."""
+    got = transformation_utils.add_op_code(
+        op_code=op_code, model_op_codes=self.model.operatorCodes
+    )
+    self.assertEqual(expected, got)
+
+  @parameterized.named_parameters(
+      dict(
+          testcase_name="float32",
+          tensor_data=np.array([1.0, 2.0, 3.0, 4.0], dtype=np.float32),
+          tensor_type=schema_py_generated.TensorType.FLOAT32,
+          expected_type=schema_py_generated.TensorType.FLOAT32,
+          expected_shape=(4,),
+          expected_buffer_data=np.frombuffer(
+              np.array([1.0, 2.0, 3.0, 4.0], dtype=np.float32).tobytes(),
+              dtype=np.uint8,
+          ).flatten(),
+      ),
+      dict(
+          testcase_name="int8",
+          tensor_data=np.array([[1, 2], [3, 4]], dtype=np.int8),
+          tensor_type=schema_py_generated.TensorType.INT8,
+          expected_type=schema_py_generated.TensorType.INT8,
+          expected_shape=(2, 2),
+          expected_buffer_data=np.frombuffer(
+              np.array([[1, 2], [3, 4]], dtype=np.int8).tobytes(),
+              dtype=np.uint8,
+          ).flatten(),
+      ),
+  )
+  def test_add_new_constant_tensor(
+      self,
+      tensor_data,
+      tensor_type,
+      expected_type,
+      expected_shape,
+      expected_buffer_data,
+  ):
+    """Tests if the constant tensor is added to the model."""
+    ret = transformation_utils.add_new_constant_tensor(
+        tensor_name="test_tensor",
+        data=tensor_data,
+        tensor_type=tensor_type,
+        subgraph=self.model.subgraphs[0],
+        buffers=self.model.buffers,
+    )
+    self.assertEqual(ret, len(self.model.subgraphs[0].tensors) - 1)
+    self.assertEqual(
+        str(self.model.subgraphs[0].tensors[-1].name), "test_tensor"
+    )
+    self.assertEqual(
+        expected_type,
+        self.model.subgraphs[0].tensors[-1].type,
+    )
+    self.assertEqual(
+        expected_shape,
+        self.model.subgraphs[0].tensors[-1].shape,
+    )
+    self.assertListEqual(
+        expected_buffer_data.tolist(),
+        self.model.buffers[
+            self.model.subgraphs[0].tensors[-1].buffer
+        ].data.tolist(),
+    )
+
+  @parameterized.named_parameters(
+      dict(
+          testcase_name="float32",
+          tensor_type=schema_py_generated.TensorType.FLOAT32,
+          tensor_shape=[1, 1, 1, 1],
+          expected_shape=[1, 1, 1, 1],
+          expected_type=schema_py_generated.TensorType.FLOAT32,
+      ),
+      dict(
+          testcase_name="int8",
+          tensor_type=schema_py_generated.TensorType.INT8,
+          tensor_shape=[1, 224, 224, 1],
+          expected_shape=[1, 224, 224, 1],
+          expected_type=schema_py_generated.TensorType.INT8,
+      ),
+  )
+  def test_add_new_activation_tensor_to_subgraph(
+      self,
+      tensor_type,
+      tensor_shape,
+      expected_shape,
+      expected_type,
+  ):
+    """Tests if the activation tensor is added to the subgraph."""
+    ret = transformation_utils.add_new_activation_tensor(
+        tensor_name="test_tensor",
+        shape=tensor_shape,
+        tensor_type=tensor_type,
+        subgraph=self.model.subgraphs[0],
+    )
+    self.assertEqual(ret, len(self.model.subgraphs[0].tensors) - 1)
+    self.assertEqual(
+        str(self.model.subgraphs[0].tensors[-1].name), "test_tensor"
+    )
+    self.assertEqual(
+        expected_type,
+        self.model.subgraphs[0].tensors[-1].type,
+    )
+    self.assertEqual(
+        expected_shape,
+        self.model.subgraphs[0].tensors[-1].shape,
+    )
+
+
+if __name__ == "__main__":
+  googletest.main()

ai_edge_quantizer/utils/__init__.py

@@ -0,0 +1,15 @@
+# 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.
+# ==============================================================================
+

ai_edge_quantizer/utils/calibration_utils.py

@@ -0,0 +1,86 @@
+# 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.
+# ==============================================================================
+
+"""Utilities for model calibration."""
+
+from typing import Union
+import numpy as np
+from ai_edge_quantizer import qtyping
+
+
+def _update_moving_average(
+    smoothing_factor: Union[np.ndarray, float],
+    w: np.ndarray,
+    update: np.ndarray,
+) -> np.ndarray:
+  """Updates weight w with moving average.
+
+  Args:
+    smoothing_factor: Smoothing factor used to update w.
+    w: Weights to be updated.
+    update: Value used for update.
+
+  Returns:
+    Weighted sum of w and update.
+  """
+  return smoothing_factor * w + (1.0 - smoothing_factor) * update
+
+
+def moving_average_update(
+    qsv: qtyping.QSV, new_qsv: qtyping.QSV, smoothing_factor: float = 0.95
+) -> qtyping.QSV:
+  """Update the QSV (i.e., min/max) using moving average.
+
+  Args:
+    qsv: The quantization statistical value of the tensor (min/max) that need to
+      be updated.
+    new_qsv: The new QSVs (e.g., from new round of calibration).
+    smoothing_factor: The weight of moving average.
+
+  Returns:
+    The updated QSV for the tensor.
+  """
+  if not qsv:
+    return new_qsv
+
+  updated_qsv = {}
+  updated_qsv["min"] = _update_moving_average(
+      smoothing_factor, qsv["min"], new_qsv["min"]
+  )
+
+  updated_qsv["max"] = _update_moving_average(
+      smoothing_factor, qsv["max"], new_qsv["max"]
+  )
+  return updated_qsv
+
+
+def min_max_update(qsv: qtyping.QSV, new_qsv: qtyping.QSV) -> qtyping.QSV:
+  """Update the QSV with minimum min values and maximum max values.
+
+  Args:
+    qsv: The quantization statistical value of the tensor (min/max) that need to
+      be updated.
+    new_qsv: The new QSVs (e.g., from new round of calibration).
+
+  Returns:
+    The updated QSV for the tensor.
+  """
+  if not qsv:
+    return new_qsv
+
+  updated_qsv = {}
+  updated_qsv["min"] = np.minimum(qsv["min"], new_qsv["min"])
+  updated_qsv["max"] = np.maximum(qsv["max"], new_qsv["max"])
+  return updated_qsv

ai_edge_quantizer/utils/calibration_utils_test.py

@@ -0,0 +1,77 @@
+# 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 absl.testing import parameterized
+from tensorflow.python.platform import googletest
+from ai_edge_quantizer.utils import calibration_utils
+
+
+class CalibrationUtilsTest(parameterized.TestCase):
+
+  @parameterized.named_parameters(
+      dict(
+          testcase_name="zero_smoothing_factor",
+          smoothing_factor=0,
+          expected_vals={"min": -1000, "max": 800},
+      ),
+      dict(
+          testcase_name="one_smoothing_factor",
+          smoothing_factor=1,
+          expected_vals={"min": -10, "max": 8},
+      ),
+      dict(
+          testcase_name="normal_smoothing_factor",
+          smoothing_factor=0.99,
+          expected_vals={"min": -19.9, "max": 15.92},
+      ),
+  )
+  def test_update_tensor_qsv_moving_average(
+      self, smoothing_factor, expected_vals
+  ):
+    old_qsv = {"min": -10, "max": 8}
+    # Large values to mimic outlier.
+    new_qsv = {"min": -1000, "max": 800}
+    updated_qsv = calibration_utils.moving_average_update(
+        old_qsv, new_qsv, smoothing_factor=smoothing_factor
+    )
+    self.assertAlmostEqual(updated_qsv["min"], expected_vals["min"])
+    self.assertAlmostEqual(updated_qsv["max"], expected_vals["max"])
+
+  @parameterized.named_parameters(
+      dict(
+          testcase_name="scalar",
+          old_qsv={"min": -10, "max": 8},
+          new_qsv={"min": -1000, "max": 1},
+          expected_qsv={"min": -1000, "max": 8},
+      ),
+      dict(
+          testcase_name="2darray",
+          old_qsv={"min": [[-19], [20]], "max": [[21], [250]]},
+          new_qsv={"min": [[-1000], [25]], "max": [[33], [100]]},
+          expected_qsv={"min": [[-1000], [20]], "max": [[33], [250]]},
+      ),
+  )
+  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"])
+    else:
+      self.assertEqual(updated_qsv["min"], expected_qsv["min"])
+      self.assertEqual(updated_qsv["max"], expected_qsv["max"])
+
+
+if __name__ == "__main__":
+  googletest.main()

ai_edge_quantizer/utils/test_utils.py

@@ -0,0 +1,107 @@
+# 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 tests."""
+
+import inspect as _inspect
+import os.path as _os_path
+import sys as _sys
+from typing import Any, Union
+
+import numpy as np
+
+from ai_edge_quantizer.utils import tfl_interpreter_utils
+
+
+def get_path_to_datafile(path):
+  """Get the path to the specified file in the data dependencies.
+
+  The path is relative to the file calling the function.
+
+  Args:
+    path: a string resource path relative to the calling file.
+
+  Returns:
+    The path to the specified file present in the data attribute of py_test
+    or py_binary.
+
+  Raises:
+    IOError: If the path is not found, or the resource can't be opened.
+  """
+  data_files_path = _os_path.dirname(_inspect.getfile(_sys._getframe(1)))  # pylint: disable=protected-access
+  path = _os_path.join(data_files_path, path)
+  path = _os_path.normpath(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']
+      )
+      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
+    )
+  return test_data