quantizer.py

# 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 API."""

from collections.abc import Iterable
import dataclasses
import json
import os
from typing import Any, Optional, Union
from ai_edge_quantizer import algorithm_manager
from ai_edge_quantizer import calibrator
from ai_edge_quantizer import default_policy
from ai_edge_quantizer import model_modifier
from ai_edge_quantizer import model_validator
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 ai_edge_quantizer.utils import tfl_interpreter_utils
from ai_edge_quantizer.utils import validation_utils
from tensorflow.python.platform import gfile  # pylint: disable=g-direct-tensorflow-import

# Expose algorithm names to users.
AlgorithmName = algorithm_manager.AlgorithmName

_QuantRecipe = recipe_manager.ModelQuantizationRecipe
_TFLOpName = qtyping.TFLOperationName
_OpQuantizationConfig = qtyping.OpQuantizationConfig
_TensorQuantizationConfig = qtyping.TensorQuantizationConfig
_TensorTransformationParams = dict[str, qtyping.TensorTransformationParams]
_SignatureInput = dict[str, Any]  # input_argument_name -> tensor_value.
_CalibrationResult = dict[str, qtyping.QSV]


@dataclasses.dataclass(frozen=True)
class QuantizationResult:
  """Quantization result.

  Attributes:
    recipe: Quantization recipe.
    quantized_model: Quantized model.
  """

  recipe: _QuantRecipe
  quantized_model: Optional[bytearray]

  def save(self, save_folder: str, model_name: str) -> None:
    """Saves the quantized model and the quantization recipe.

    Args:
      save_folder: Path to the folder to save the quantized model and the
        quantization recipe.
      model_name: Name of the model.

    Raises:
      RuntimeError: If no quantized model is available.
      FileExistsError: If the model already exists in the folder.
    """
    if self.quantized_model is None:
      raise RuntimeError(
          'No quantized model to save. Make sure .quantize() is called.'
      )
    model_save_path = os.path.join(save_folder, f'{model_name}.tflite')
    if gfile.Exists(model_save_path):
      raise FileExistsError(
          f'The model {model_save_path} already exists in the folder.'
      )
    with gfile.GFile(model_save_path, 'wb') as output_file_handle:
      output_file_handle.write(self.quantized_model)

    recipe = json.dumps(self.recipe)
    recipe_save_path = os.path.join(save_folder, model_name + '_recipe.json')
    with gfile.GFile(recipe_save_path, 'w') as output_file_handle:
      output_file_handle.write(recipe)

  def export_model(self, filepath: str) -> None:
    """Exports the quantized model to a .tflite flatbuffer.

    Args:
      filepath: Path (including file name) that the exported model should be
        serialized to.

    Raises:
      RuntimeError: If no quantized model is available.
    """
    if self.quantized_model is None:
      raise RuntimeError(
          'No quantized model to save. Make sure .quantize() is called.'
      )
    with gfile.GFile(filepath, 'wb') as output_file_handle:
      output_file_handle.write(self.quantized_model)


class Quantizer:
  """AI Edge Quantizer API.

  Attributes:
    float_model: TFLite model file path or bytearray.
    quantization_recipe: Quantization recipe .json filepath or in loaded json
      format.
  """

  def __init__(
      self,
      float_model: Union[str, bytearray],
      quantization_recipe: Optional[Union[str, _QuantRecipe]] = None,
  ):
    """Initializes the quantizer.

    Args:
      float_model: Path to the float tflite model.
      quantization_recipe: Quantization recipe in .json filepath or loaded json
        format.
    """
    # Use `float model` as bytes for memory efficiency.
    self.float_model: bytes = (
        tfl_flatbuffer_utils.get_model_content(float_model)
        if isinstance(float_model, str)
        else float_model
    )

    self._recipe_manager: recipe_manager.RecipeManager = (
        recipe_manager.RecipeManager()
    )
    if quantization_recipe is not None:
      self.load_quantization_recipe(quantization_recipe)
    self._result: QuantizationResult = QuantizationResult([{}], None)

  def load_quantization_recipe(self, recipe: Union[str, _QuantRecipe]) -> None:
    """Loads a quantization recipe.

    The existing recipe will be overwritten.

    Args:
      recipe: Quantization recipe in json format.
    """
    if isinstance(recipe, str):
      with gfile.Open(recipe) as json_file:
        recipe = json.load(json_file)
    self._recipe_manager.load_quantization_recipe(recipe)

  def load_config_policy(self, filename: str) -> None:
    """Loads a JSON policy.

    The existing policy will be overwritten.

    Args:
      filename: Config policy filename.
    """
    with gfile.Open(filename, 'r') as f:
      policy = default_policy.update_default_config_policy(f.read())

    # Register the policy for MIN_MAX_UNIFORM_QUANT algorithm.
    algorithm_manager.register_config_check_policy_func(
        AlgorithmName.MIN_MAX_UNIFORM_QUANT, policy
    )

  def get_quantization_recipe(self) -> _QuantRecipe:
    """Gets the quantization recipe.

    Returns:
      A quantization recipe.
    """
    return self._recipe_manager.get_quantization_recipe()

  def update_quantization_recipe(
      self,
      regex: str,
      operation_name: _TFLOpName,
      op_config: Optional[_OpQuantizationConfig] = None,
      algorithm_key: str = algorithm_manager.AlgorithmName.MIN_MAX_UNIFORM_QUANT,
  ):
    """Adds a quantization configuration to the recipe.

    Conflict arises when we are trying to set an operation under a certain regex
    which is already existed in the config dictionary. Under such circumstance,
    the new config is used to replace the previous one.

    We also have special treatment for _TFLOperationKey.ALL. If the new config
    is on _TFLOperationKey.ALL and there are existing op configs inside the same
    scope, we clear the previous configs and use _TFLOperationKey.ALL.

    Args:
      regex: Regular expression for layer name matching.
      operation_name: Target TFLite operation. * for all supported TFLite
        operation.
      op_config: Quantization configuration which will be used to update the
        default configuration. None or empty dict means the default
        configuration will be used.
      algorithm_key: Algorithm key to be applied.
    """
    self._recipe_manager.add_quantization_config(
        regex, operation_name, op_config, algorithm_key
    )

  @property
  def need_calibration(self) -> bool:
    """Checks if the current recipe needs calibration."""
    return self._recipe_manager.need_calibration()

  def calibrate(
      self,
      calibration_data: dict[str, Iterable[_SignatureInput]],
      previous_calibration_result: Optional[_CalibrationResult] = None,
      num_threads: int = 16,
  ) -> _CalibrationResult:
    """Calibrates the float model (required by static range quantization).

    Args:
      calibration_data: Calibration data for a model signature.
      previous_calibration_result: Previous calibration result to be loaded. The
        calibration process will be resumed from the previous result.
      num_threads: Number of threads to use for calibration.

    Returns:
      Calibration result ({tensor_name: tensor QSVs (e.g.,min/max)}).

    Raises:
      ValueError: If the calibration result is insufficient.
    """
    if not self.need_calibration:
      return {}

    calib = calibrator.Calibrator(self.float_model, num_threads=num_threads)
    if previous_calibration_result is not None:
      calib.load_model_qsvs(previous_calibration_result)
    calib.calibrate(calibration_data, self._recipe_manager)
    return calib.get_model_qsvs()

  def _ensure_model_qsv_sufficient(
      self, calibration_result: _CalibrationResult
  ):
    """Checks if the calibration result has sufficient QSV."""

    # Find all tensor names with empty entries.
    empty_qsvs = [key for key, value in calibration_result.items() if not value]

    # Go over every signature and check if empty entry tensor belongs to it.
    tfl_interpreter = tfl_interpreter_utils.create_tfl_interpreter(
        self.float_model
    )
    for signature_key in tfl_interpreter.get_signature_list():
      subgraph_idx = tfl_interpreter_utils.get_signature_main_subgraph_index(
          tfl_interpreter, signature_key
      )

      for tensor_detail in tfl_interpreter.get_tensor_details(subgraph_idx):
        tensor_name = tensor_detail['name']
        if tensor_name in empty_qsvs:
          raise ValueError(
              f'Missing QSVs (min/max) for tensor {tensor_name} in Signature'
              f" '{signature_key}'. Please check if Signature"
              f' {signature_key} has been calibrated.'
          )

  def quantize(
      self, calibration_result: Optional[_CalibrationResult] = None
  ) -> QuantizationResult:
    """Quantizes the float model.

    Args:
      calibration_result: Calibration result to be used for quantization (if
        needed, check with self.need_calibration).

    Returns:
      Quantization result.

    Raises:
      RuntimeError: If quantization recipe is empty.
    """

    if calibration_result is not None:
      self._ensure_model_qsv_sufficient(calibration_result)

    if not self.get_quantization_recipe():
      raise RuntimeError('Can not quantize without a quantization recipe.')
    quant_params = self._get_quantization_params(calibration_result)
    quantized_model = self._get_quantized_model(quant_params)
    self._result = QuantizationResult(
        self.get_quantization_recipe(), quantized_model
    )
    return self._result

  def validate(
      self,
      test_data: Optional[dict[str, Iterable[_SignatureInput]]] = None,
      error_metrics: str = 'mse',
      use_xnnpack: bool = True,
      num_threads: int = 16,
  ) -> model_validator.ComparisonResult:
    """Numerical validation of the quantized model for a model signature.

    Side by side numerical comparison will be performed on all tensors in the
    quantized model against ones from the float model. If no test data is
    provided, random normal distributed data will be used. This test is intended
    to be SANITY check for the quality of the quantized model. End to end task
    specific test should be performed as the golden standard of the quantized
    model quality. The comparison result will be saved in json format if
    json_save_path is provided.

    Args:
      test_data: A dictionary of signature key and its correspending test input
        data that will be used for validation. If set to None, random normal
        distributed data will be used for all signatures in the model.
      error_metrics: Error metrics to be used for comparison.
      use_xnnpack: Whether to use the xnnpack library for validation.
      num_threads: Number of threads to use for validation.

    Returns:
      The comparison result.
    """
    if test_data is None:
      # Create test data for all signatures in the model.
      test_data = test_utils.create_random_normal_input_data(
          self.float_model, num_samples=1
      )
    return model_validator.compare_model(
        self.float_model,
        self._result.quantized_model,
        test_data,
        error_metrics,
        validation_utils.get_validation_func(error_metrics),
        use_xnnpack=use_xnnpack,
        num_threads=num_threads,
    )

  def _get_quantization_params(
      self, calibration_result: Optional[_CalibrationResult] = None
  ) -> _TensorTransformationParams:
    """Gets the quantization parameters.

    Args:
      calibration_result: Calibration result to be used for quantization (if
        needed, check with self.need_calibration).

    Returns:
      A dictionary containing the quantization parameters.
    """
    params_generator_instance = params_generator.ParamsGenerator(
        self.float_model
    )
    return params_generator_instance.generate_quantization_parameters(
        self._recipe_manager, calibration_result
    )

  def _get_quantized_model(
      self, quant_params: _TensorTransformationParams
  ) -> bytearray:
    """Gets the quantized model.

    Args:
      quant_params: A dictionary containing the quantization parameters.

    Returns:
      The quantized model.
    """
    model_modifier_instance = model_modifier.ModelModifier(self.float_model)
    return model_modifier_instance.modify_model(quant_params)