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

ai_edge_quantizer/default_policy.py

@@ -0,0 +1,310 @@
+# 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.
+# ==============================================================================
+
+"""Default quantization policy."""
+
+import collections
+import copy
+import json
+from typing import Any
+from ai_edge_quantizer import qtyping
+
+_TFLOpName = qtyping.TFLOperationName
+_OpQuantizationConfig = qtyping.OpQuantizationConfig
+_TensorQuantizationConfig = qtyping.TensorQuantizationConfig
+_ComputePrecision = qtyping.ComputePrecision
+_Granularity = qtyping.QuantGranularity
+_INT = qtyping.TensorDataType.INT
+
+# Default config check policy in JSON format. It has two keys: "configs" and
+# "ops_per_config". "configs" is a dictionary mapping from config name to config
+# content. "ops_per_config" is a dictionary mapping from config name to a list
+# of op names.
+DEFAULT_JSON_POLICY = """
+{
+  "configs": {
+    "dynamic_wi8_afp32": {
+      "weight_tensor_config": {
+        "num_bits": 8,
+        "symmetric": [true],
+        "granularity": ["CHANNELWISE", "TENSORWISE"],
+        "dtype": "INT"
+      },
+      "explicit_dequantize": false,
+      "compute_precision": "INTEGER"
+    },
+    "dynamic_wi4_afp32": {
+      "weight_tensor_config": {
+        "num_bits": 4,
+        "symmetric": [true],
+        "granularity": ["CHANNELWISE", "TENSORWISE"],
+        "dtype": "INT"
+      },
+      "explicit_dequantize": false,
+      "compute_precision": "INTEGER"
+    },
+    "static_wi8_ai16": {
+      "activation_tensor_config": {
+        "num_bits": 16,
+        "symmetric": [true],
+        "granularity": ["TENSORWISE"],
+        "dtype": "INT"
+      },
+      "weight_tensor_config": {
+        "num_bits": 8,
+        "symmetric": [true],
+        "granularity": ["CHANNELWISE", "TENSORWISE"],
+        "dtype": "INT"
+      },
+      "explicit_dequantize": false,
+      "compute_precision": "INTEGER"
+    },
+    "static_wi4_ai16": {
+      "activation_tensor_config": {
+        "num_bits": 16,
+        "symmetric": [true],
+        "granularity": ["TENSORWISE"],
+        "dtype": "INT"
+      },
+      "weight_tensor_config": {
+        "num_bits": 4,
+        "symmetric": [true],
+        "granularity": ["CHANNELWISE", "TENSORWISE"],
+        "dtype": "INT"
+      },
+      "explicit_dequantize": false,
+      "compute_precision": "INTEGER"
+    },
+    "static_wi8_ai8": {
+      "activation_tensor_config": {
+        "num_bits": 8,
+        "symmetric": [true, false],
+        "granularity": ["TENSORWISE"],
+        "dtype": "INT"
+      },
+      "weight_tensor_config": {
+        "num_bits": 8,
+        "symmetric": [true],
+        "granularity": ["CHANNELWISE", "TENSORWISE"],
+        "dtype": "INT"
+      },
+      "explicit_dequantize": false,
+      "compute_precision": "INTEGER"
+    },
+    "static_wi4_ai8": {
+      "activation_tensor_config": {
+        "num_bits": 8,
+        "symmetric": [true, false],
+        "granularity": ["TENSORWISE"],
+        "dtype": "INT"
+      },
+      "weight_tensor_config": {
+        "num_bits": 4,
+        "symmetric": [true],
+        "granularity": ["CHANNELWISE", "TENSORWISE"],
+        "dtype": "INT"
+      },
+      "explicit_dequantize": false,
+      "compute_precision": "INTEGER"
+    },
+    "weightonly_wi8_afp32": {
+      "weight_tensor_config": {
+        "num_bits": 8,
+        "symmetric": [true, false],
+        "granularity": ["CHANNELWISE", "TENSORWISE"],
+        "dtype": "INT"
+      },
+      "explicit_dequantize": true,
+      "compute_precision": "FLOAT"
+    },
+    "weightonly_wi4_afp32": {
+      "weight_tensor_config": {
+        "num_bits": 4,
+        "symmetric": [true, false],
+        "granularity": ["CHANNELWISE", "TENSORWISE"],
+        "dtype": "INT"
+      },
+      "explicit_dequantize": true,
+      "compute_precision": "FLOAT"
+    }
+  },
+  "ops_per_config": {
+    "static_wi8_ai16": [
+      "ADD",
+      "AVERAGE_POOL_2D",
+      "BATCH_MATMUL",
+      "CONCATENATION",
+      "CONV_2D",
+      "CONV_2D_TRANSPOSE",
+      "DEPTHWISE_CONV_2D",
+      "FULLY_CONNECTED",
+      "GELU",
+      "LOGISTIC",
+      "MEAN",
+      "MUL",
+      "RESHAPE",
+      "RSQRT",
+      "SOFTMAX",
+      "SPLIT",
+      "STRIDED_SLICE",
+      "SUB",
+      "TANH",
+      "TRANSPOSE",
+      "INPUT",
+      "OUTPUT",
+      "SLICE",
+      "EMBEDDING_LOOKUP",
+      "SUM",
+      "SELECT_V2"
+    ],
+    "static_wi8_ai8": [
+      "ADD",
+      "AVERAGE_POOL_2D",
+      "BATCH_MATMUL",
+      "CONCATENATION",
+      "FULLY_CONNECTED",
+      "CONV_2D",
+      "CONV_2D_TRANSPOSE",
+      "DEPTHWISE_CONV_2D",
+      "GELU",
+      "LOGISTIC",
+      "MEAN",
+      "MUL",
+      "RESHAPE",
+      "RSQRT",
+      "SOFTMAX",
+      "SPLIT",
+      "STRIDED_SLICE",
+      "SUB",
+      "TANH",
+      "TRANSPOSE",
+      "INPUT",
+      "OUTPUT",
+      "SLICE",
+      "EMBEDDING_LOOKUP",
+      "SUM",
+      "SELECT_V2"
+    ],
+    "static_wi4_ai8": ["FULLY_CONNECTED", "CONV_2D", "INPUT", "OUTPUT"],
+    "static_wi4_ai16": ["FULLY_CONNECTED", "CONV_2D", "INPUT", "OUTPUT"],
+    "dynamic_wi8_afp32": [
+      "BATCH_MATMUL",
+      "CONV_2D",
+      "CONV_2D_TRANSPOSE",
+      "DEPTHWISE_CONV_2D",
+      "EMBEDDING_LOOKUP",
+      "FULLY_CONNECTED"
+    ],
+    "dynamic_wi4_afp32": ["FULLY_CONNECTED", "EMBEDDING_LOOKUP", "CONV_2D"],
+    "weightonly_wi8_afp32": [
+      "BATCH_MATMUL",
+      "CONV_2D",
+      "CONV_2D_TRANSPOSE",
+      "DEPTHWISE_CONV_2D",
+      "EMBEDDING_LOOKUP",
+      "FULLY_CONNECTED"
+    ],
+    "weightonly_wi4_afp32": ["BATCH_MATMUL", "FULLY_CONNECTED", "EMBEDDING_LOOKUP", "CONV_2D"]
+  }
+}
+"""
+
+
+def _unroll_json_config(
+    json_config: dict[str, Any],
+) -> list[_OpQuantizationConfig]:
+  """Unrolls the config.
+
+  Args:
+    json_config: JSON config to be unrolled.
+
+  Returns:
+    quant_configs: A list of unrolled configs.
+  """
+
+  # Unroll activation configs first.
+  activation_configs = []
+  if "activation_tensor_config" in json_config:
+    for symmetric in json_config["activation_tensor_config"]["symmetric"]:
+      for granularity in json_config["activation_tensor_config"]["granularity"]:
+        tensor_config = {
+            "num_bits": json_config["activation_tensor_config"]["num_bits"],
+            "symmetric": symmetric,
+            "granularity": granularity,
+            "dtype": json_config["activation_tensor_config"]["dtype"],
+        }
+        activation_configs.append(
+            qtyping.TensorQuantizationConfig.from_dict(tensor_config)
+        )
+
+  # Then unroll weight configs and turn them into quantization configs.
+  quant_configs = []
+  for symmetric in json_config["weight_tensor_config"]["symmetric"]:
+    for granularity in json_config["weight_tensor_config"]["granularity"]:
+      tensor_config = {
+          "num_bits": json_config["weight_tensor_config"]["num_bits"],
+          "symmetric": symmetric,
+          "granularity": granularity,
+          "dtype": json_config["weight_tensor_config"]["dtype"],
+      }
+
+      if activation_configs:
+        for activation_config in activation_configs:
+          quant_configs.append(
+              qtyping.OpQuantizationConfig(
+                  activation_tensor_config=activation_config,
+                  weight_tensor_config=qtyping.TensorQuantizationConfig.from_dict(
+                      tensor_config
+                  ),
+                  compute_precision=json_config["compute_precision"],
+                  explicit_dequantize=json_config["explicit_dequantize"],
+              )
+          )
+      else:
+        quant_configs.append(
+            qtyping.OpQuantizationConfig(
+                weight_tensor_config=qtyping.TensorQuantizationConfig.from_dict(
+                    tensor_config
+                ),
+                compute_precision=json_config["compute_precision"],
+                explicit_dequantize=json_config["explicit_dequantize"],
+            )
+        )
+
+  return quant_configs
+
+
+def update_default_config_policy(raw_json_policy: str):
+  """Updates the default config check policy."""
+  json_policy_content = json.loads(raw_json_policy)
+
+  # Unroll the json config and add the configs to the policy.
+  policy = collections.OrderedDict()
+  for json_policy_config in json_policy_content["ops_per_config"]:
+    unrolled_configs = _unroll_json_config(
+        json_policy_content["configs"][json_policy_config]
+    )
+
+    for op in json_policy_content["ops_per_config"][json_policy_config]:
+      op_name = _TFLOpName(op)
+      quant_configs = copy.deepcopy(unrolled_configs)
+      if op in policy.keys():
+        quant_configs += policy[op_name]
+      policy[op_name] = quant_configs
+
+  return policy
+
+
+DEFAULT_CONFIG_CHECK_POLICY = update_default_config_policy(DEFAULT_JSON_POLICY)

ai_edge_quantizer/model_modifier.py

@@ -0,0 +1,176 @@
+# 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.
+# ==============================================================================
+
+"""Model Modifier class that produce the final quantized TFlite model."""
+
+import copy
+
+import numpy as np
+
+from ai_edge_quantizer import qtyping
+from ai_edge_quantizer import transformation_instruction_generator
+from ai_edge_quantizer import transformation_performer
+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
+
+
+class ModelModifier:
+  """Model Modifier class that produce the final quantized TFlite model."""
+
+  def __init__(self, float_tflite: bytes):
+    """Constructor.
+
+    Args:
+      float_tflite: the original TFlite model in bytearray or file path
+    """
+
+    self._model_content = float_tflite
+
+    self._constant_map = []
+    self._transformation_instruction_generator = (
+        transformation_instruction_generator.TransformationInstructionsGenerator()
+    )
+    self._transformation_performer = (
+        transformation_performer.TransformationPerformer()
+    )
+
+  def modify_model(
+      self, params: dict[str, qtyping.TensorTransformationParams]
+  ) -> bytearray:
+    """Modify the model.
+
+    Args:
+      params: a dictionary with tensor name and a list of tensor transformation
+        params
+
+    Returns:
+      a byte buffer that represents the serialized tflite model
+    """
+    quantized_model = copy.deepcopy(
+        flatbuffer_utils.read_model_from_bytearray(self._model_content)
+    )
+
+    instructions = self._transformation_instruction_generator.quant_params_to_transformation_insts(
+        params, quantized_model
+    )
+
+    self._transformation_performer.transform_graph(
+        instructions, quantized_model
+    )
+    constant_buffer_size = self._process_constant_map(quantized_model)
+    # we leave 64MB for the model architecture.
+    if constant_buffer_size > 2**31 - 2**26:
+      return self._serialize_large_model(quantized_model)
+    else:
+      return self._serialize_small_model(quantized_model)
+
+  def _process_constant_map(
+      self, quantized_model: schema_py_generated.ModelT
+  ) -> int:
+    """Process the constant map after all transformations are applied.
+
+    If the resulting model is > 2GB then we would need to serialize constants
+    separately, as such, we collect all the constant buffers using this
+    function.
+
+    Args:
+      quantized_model: a quantized TFlite ModelT
+
+    Returns:
+      an integer representing the total size of the constant buffers
+    """
+    buffer_size = 0
+    for buffer in quantized_model.buffers:
+      if buffer.data is None:
+        self._constant_map.append(buffer.data)
+      elif isinstance(buffer.data, np.ndarray):
+        self._constant_map.append(buffer.data.tobytes())
+        buffer_size += len(buffer.data.tobytes())
+      else:
+        self._constant_map.append(buffer.data)
+        buffer_size += len(buffer.data)
+    return buffer_size
+
+  def _pad_bytearray(self, bytearr: bytearray):
+    """Pad the bytearray to 16 bytes."""
+    remainder = len(bytearr) % 16
+    if remainder != 0:
+      padding_size = 16 - remainder
+      bytearr.extend(b'\0' * padding_size)
+
+  # TODO: b/333797307 - support > 2GB output model
+  def _serialize_large_model(
+      self, quantized_model: schema_py_generated.ModelT
+  ) -> bytearray:
+    """serialize models > 2GB.
+
+    Args:
+      quantized_model: a quantized TFlite ModelT
+
+    Returns:
+      a byte buffer that represents the serialized tflite model
+    """
+    # TODO: b/338244867 - we can have more efficient way to calculate the
+    # buffer offsets.
+
+    # remove all the constant from the model.
+    for buffer in quantized_model.buffers:
+      if buffer.data is not None:
+        buffer.data = None
+        buffer.offset = 1
+        buffer.size = 1
+    dummy_bytearray = bytearray(
+        flatbuffer_utils.convert_object_to_bytearray(quantized_model)
+    )
+    # calculate the correct buffer size and offset
+    self._pad_bytearray(dummy_bytearray)
+    for buffer_idx, buffer in enumerate(quantized_model.buffers):
+      buffer_data = self._constant_map[buffer_idx]
+      if buffer_data is None:
+        continue
+      buffer.offset = len(dummy_bytearray)
+      buffer.size = len(buffer_data)
+      dummy_bytearray += buffer_data
+      self._pad_bytearray(dummy_bytearray)
+    del dummy_bytearray
+
+    # build new tflite file with correct buffer offset
+    model_bytearray = bytearray(
+        flatbuffer_utils.convert_object_to_bytearray(quantized_model)
+    )
+    self._pad_bytearray(model_bytearray)
+    for buffer_idx, _ in enumerate(quantized_model.buffers):
+      buffer_data = self._constant_map[buffer_idx]
+      if buffer_data is None:
+        continue
+      model_bytearray += buffer_data
+      self._pad_bytearray(model_bytearray)
+    return model_bytearray
+
+  def _serialize_small_model(
+      self, quantized_model: schema_py_generated.ModelT
+  ) -> bytearray:
+    """serialize models < 2GB.
+
+    Args:
+      quantized_model: a quantized TFlite ModelT
+
+    Returns:
+      a byte buffer that represents the serialized tflite model
+    """
+    model_bytearray = flatbuffer_utils.convert_object_to_bytearray(
+        quantized_model
+    )
+    return model_bytearray

ai_edge_quantizer/model_modifier_test.py

@@ -0,0 +1,130 @@
+# 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 model_modifier."""
+
+import os
+import tracemalloc
+from tensorflow.python.platform import googletest
+from absl.testing import parameterized
+from ai_edge_quantizer import model_modifier
+from ai_edge_quantizer import params_generator
+from ai_edge_quantizer import qtyping
+from ai_edge_quantizer import recipe_manager
+from ai_edge_quantizer.utils import test_utils
+from ai_edge_quantizer.utils import tfl_flatbuffer_utils
+from tensorflow.lite.tools import flatbuffer_utils  # pylint: disable=g-direct-tensorflow-import
+
+TEST_DATA_PREFIX_PATH = test_utils.get_path_to_datafile('.')
+
+
+class ModelModifierTest(parameterized.TestCase):
+
+  def setUp(self):
+    super().setUp()
+    self._model_path = os.path.join(
+        TEST_DATA_PREFIX_PATH, 'tests/models/conv_fc_mnist.tflite'
+    )
+
+    self._model_content: bytes = tfl_flatbuffer_utils.get_model_content(
+        self._model_path
+    )
+    self._model_modifier = model_modifier.ModelModifier(self._model_content)
+    self._global_recipe = [
+        {
+            'regex': '.*',
+            'operation': 'FULLY_CONNECTED',
+            'algorithm_key': 'min_max_uniform_quantize',
+            'op_config': {
+                'weight_tensor_config': {
+                    'dtype': qtyping.TensorDataType.INT,
+                    'num_bits': 8,
+                    'symmetric': False,
+                    'granularity': qtyping.QuantGranularity.CHANNELWISE,
+                    'block_size': 0,
+                },
+                # Equivalent to WEIGHT_ONLY.
+                'compute_precision': qtyping.ComputePrecision.FLOAT,
+                'explicit_dequantize': True,
+            },
+        },
+    ]
+
+  def test_process_constant_map_succeeds(self):
+    model_bytearray = flatbuffer_utils.read_model_from_bytearray(
+        self._model_content
+    )
+    constant_size = self._model_modifier._process_constant_map(model_bytearray)
+    self.assertEqual(constant_size, 202540)
+
+  def test_modify_model_succeeds_with_recipe(self):
+    recipe_manager_instance = recipe_manager.RecipeManager()
+    params_generator_instance = params_generator.ParamsGenerator(
+        self._model_path
+    )
+
+    recipe_manager_instance.load_quantization_recipe(self._global_recipe)
+    tensor_quantization_params = (
+        params_generator_instance.generate_quantization_parameters(
+            recipe_manager_instance
+        )
+    )
+    new_model_binary = self._model_modifier.modify_model(
+        tensor_quantization_params
+    )
+    flatbuffer_utils.convert_bytearray_to_object(new_model_binary)
+    self.assertLess(new_model_binary, self._model_content)
+
+  def test_modify_model_preserves_original_model(self):
+    recipe_manager_instance = recipe_manager.RecipeManager()
+    params_generator_instance = params_generator.ParamsGenerator(
+        self._model_path
+    )
+
+    recipe_manager_instance.load_quantization_recipe(self._global_recipe)
+    tensor_quantization_params = (
+        params_generator_instance.generate_quantization_parameters(
+            recipe_manager_instance
+        )
+    )
+    self.assertEqual(self._model_modifier._model_content, self._model_content)
+    self._model_modifier.modify_model(tensor_quantization_params)
+    self.assertEqual(self._model_modifier._model_content, self._model_content)
+
+  def test_modify_model_peak_memory_usage_in_acceptable_range(self):
+    """Test ModifyModel peak memory usage."""
+
+    recipe_manager_instance = recipe_manager.RecipeManager()
+    params_generator_instance = params_generator.ParamsGenerator(
+        self._model_path
+    )
+
+    recipe_manager_instance.load_quantization_recipe(self._global_recipe)
+    tensor_quantization_params = (
+        params_generator_instance.generate_quantization_parameters(
+            recipe_manager_instance
+        )
+    )
+
+    tracemalloc.start()
+    self._model_modifier.modify_model(tensor_quantization_params)
+    _, mem_peak = tracemalloc.get_traced_memory()
+
+    loosen_mem_use_factor = 4.5
+    self.assertLess(mem_peak / len(self._model_content), loosen_mem_use_factor)
+
+
+if __name__ == '__main__':
+  googletest.main()