chore: import upstream snapshot with attribution

tensorflow/lite/tools/optimize/debugging/python/BUILD

@@ -0,0 +1,43 @@
+# QuantizationDebugger for TFLite accuracy tooling.
+load("//tensorflow:pytype.default.bzl", "pytype_strict_library")
+load("//tensorflow:tensorflow.bzl", "py_test")
+
+package(
+    # copybara:uncomment default_applicable_licenses = ["//tensorflow:LICENSE"],
+    licenses = ["notice"],
+)
+
+pytype_strict_library(
+    name = "debugger",
+    srcs = ["debugger.py"],
+    visibility = ["//visibility:public"],
+    deps = [
+        "//tensorflow/lite/python:convert",
+        "//tensorflow/lite/python:interpreter",
+        "//tensorflow/lite/python/metrics",
+        "//tensorflow/python/util:tf_export",
+        "//third_party/py/numpy",
+    ],
+)
+
+py_test(
+    name = "debugger_test",
+    srcs = [
+        "debugger_test.py",
+    ],
+    exec_properties = {"cpp_link.mem": "16g"},
+    strict_deps = True,
+    deps = [
+        ":debugger",
+        "@absl_py//absl/testing:parameterized",
+        #internal proto upb dep
+        "//third_party/py/numpy",
+        "//tensorflow:tensorflow_py",
+        "//tensorflow/lite/python:convert",
+        "//tensorflow/lite/python:lite",
+        "//tensorflow/lite/python/metrics",
+        "//tensorflow/python/framework:test_lib",
+        "//tensorflow/python/platform:client_testlib",
+        "//tensorflow/python/trackable:autotrackable",
+    ],
+)

tensorflow/lite/tools/optimize/debugging/python/debugger.py

@@ -0,0 +1,549 @@
+# Copyright 2021 The TensorFlow Authors. All Rights Reserved.
+#
+# 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.
+# ==============================================================================
+"""Python TF-Lite QuantizationDebugger."""
+import collections
+import csv
+import re
+from typing import (Any, Callable, Dict, IO, Iterable, List, Mapping, Optional,
+                    Sequence, Tuple)
+
+import numpy as np
+
+from tensorflow.lite.python import convert
+from tensorflow.lite.python import interpreter as _interpreter
+from tensorflow.lite.python.metrics import metrics as metrics_stub  # type: ignore
+from tensorflow.python.util import tf_export
+
+
+# TODO(b/198099651): move converter implementation out of lite.py
+TFLiteConverter = Any  # importing tf.lite creates circular dependency
+
+# Returns metrics based on difference of values for quantized/float ops.
+_DEFAULT_LAYER_DEBUG_METRICS = {
+    'num_elements': lambda diffs: diffs.size,
+    'stddev': np.std,
+    'mean_error': np.average,
+    'max_abs_error': lambda diffs: np.max(np.abs(diffs)),
+    'mean_squared_error': lambda diffs: np.average(diffs**2),
+}
+
+_NUMERIC_VERIFY_OP_NAME = 'NumericVerify'
+
+
+def _get_quant_params(
+    tensor_detail: Mapping[str, Any]) -> Optional[Tuple[float, int]]:
+  """Returns first scale and zero point from tensor detail, if present."""
+  quant_params = tensor_detail['quantization_parameters']
+  if not quant_params:
+    return None
+  if quant_params['scales'] and quant_params['zero_points']:
+    return (quant_params['scales'][0], quant_params['zero_points'][0])
+  return None
+
+
+@tf_export.tf_export('lite.experimental.QuantizationDebugOptions')
+class QuantizationDebugOptions:
+  """Debug options to set up a given QuantizationDebugger."""
+
+  def __init__(self,
+               layer_debug_metrics: Optional[Mapping[str,
+                                                     Callable[[np.ndarray],
+                                                              float]]] = None,
+               model_debug_metrics: Optional[Mapping[
+                   str, Callable[[Sequence[np.ndarray], Sequence[np.ndarray]],
+                                 float]]] = None,
+               layer_direct_compare_metrics: Optional[Mapping[str, Callable[
+                   [Sequence[np.ndarray], Sequence[np.ndarray], float, int],
+                   float]]] = None,
+               denylisted_ops: Optional[List[str]] = None,
+               denylisted_nodes: Optional[List[str]] = None,
+               fully_quantize: bool = False) -> None:
+    """Initializes debugger options.
+
+    Args:
+      layer_debug_metrics: a dict to specify layer debug functions
+        {function_name_str: function} where the function accepts result of
+          NumericVerify Op, which is value difference between float and
+          dequantized op results. The function returns single scalar value.
+      model_debug_metrics: a dict to specify model debug functions
+        {function_name_str: function} where the function accepts outputs from
+          two models, and returns single scalar value for a metric. (e.g.
+          accuracy, IoU)
+      layer_direct_compare_metrics: a dict to specify layer debug functions
+        {function_name_str: function}. The signature is different from that of
+          `layer_debug_metrics`, and this one gets passed (original float value,
+          original quantized value, scale, zero point). The function's
+          implementation is responsible for correctly dequantize the quantized
+          value to compare. Use this one when comparing diff is not enough.
+          (Note) quantized value is passed as int8, so cast to int32 is needed.
+      denylisted_ops: a list of op names which is expected to be removed from
+        quantization.
+      denylisted_nodes: a list of op's output tensor names to be removed from
+        quantization.
+      fully_quantize: Bool indicating whether to fully quantize the model.
+        Besides model body, the input/output will be quantized as well.
+        Corresponding to mlir_quantize's fully_quantize parameter.
+
+    Raises:
+      ValueError: when there are duplicate keys
+    """
+    self.layer_debug_metrics = layer_debug_metrics
+    self.model_debug_metrics = model_debug_metrics
+    self.layer_direct_compare_metrics = layer_direct_compare_metrics
+
+    keys = []
+    for metrics in [
+        layer_debug_metrics, model_debug_metrics, layer_direct_compare_metrics
+    ]:
+      if metrics is not None:
+        keys.extend(metrics.keys())
+    if len(keys) != len(set(keys)):
+      raise ValueError('Provided metrics have duplicate keys.')
+
+    self.denylisted_ops = denylisted_ops
+    self.denylisted_nodes = denylisted_nodes
+    self.fully_quantize = fully_quantize
+
+
+@tf_export.tf_export('lite.experimental.QuantizationDebugger')
+class QuantizationDebugger:
+  """Debugger for Quantized TensorFlow Lite debug mode models.
+
+  This can run the TensorFlow Lite converted models equipped with debug ops and
+  collect debug information. This debugger calculates statistics from
+  user-defined post-processing functions as well as default ones.
+  """
+
+  def __init__(self,
+               quant_debug_model_path: Optional[str] = None,
+               quant_debug_model_content: Optional[bytes] = None,
+               float_model_path: Optional[str] = None,
+               float_model_content: Optional[bytes] = None,
+               debug_dataset: Optional[Callable[
+                   [], Iterable[Sequence[np.ndarray]]]] = None,
+               debug_options: Optional[QuantizationDebugOptions] = None,
+               converter: Optional[TFLiteConverter] = None) -> None:
+    """Runs the TFLite debugging model with given debug options.
+
+    Args:
+      quant_debug_model_path: Path to the quantized debug TFLite model file.
+      quant_debug_model_content: Content of the quantized debug TFLite model.
+      float_model_path: Path to float TFLite model file.
+      float_model_content: Content of the float TFLite model.
+      debug_dataset: a factory function that returns dataset generator which is
+        used to generate input samples (list of np.ndarray) for the model. The
+        generated elements must have same types and shape as inputs to the
+        model.
+      debug_options: Debug options to debug the given model.
+      converter: Optional, use converter instead of quantized model.
+
+    Raises:
+      ValueError: If the debugger was unable to be created.
+
+    Attributes:
+      layer_statistics: results of error metrics for each NumericVerify op
+        results. in {layer_name: {metric_name: metric}} format.
+      model_statistics: results of error metrics for difference between float
+        and quantized models. in {metric_name: metric} format.
+    """
+    self._data_gen = debug_dataset
+    self._debug_options = debug_options or QuantizationDebugOptions()
+    self.converter = None
+    self.calibrated_model = None
+    self.float_model = None
+    self._float_interpreter = None
+    if converter is not None:
+      if self._debug_options.model_debug_metrics:
+        old_optimizations = converter.optimizations
+        self.converter = self._set_converter_options_for_float(converter)
+        self.float_model = self.converter.convert()
+        converter.optimizations = old_optimizations
+
+      self.converter = self._set_converter_options_for_calibration(converter)
+      self.calibrated_model = self.converter.convert()
+      # Converter should be already set up with all options
+      self._init_from_converter(
+          self._debug_options,
+          self.converter,
+          self.calibrated_model,
+          float_model=self.float_model)
+    else:
+      self._quant_interpreter = _interpreter.Interpreter(
+          quant_debug_model_path,
+          quant_debug_model_content,
+          experimental_preserve_all_tensors=(
+              self._debug_options.layer_direct_compare_metrics is not None))
+      if self._debug_options.model_debug_metrics:
+        self._float_interpreter = _interpreter.Interpreter(
+            float_model_path, float_model_content)
+    self._initialize_stats()
+
+  @property
+  def options(self) -> QuantizationDebugOptions:
+    return self._debug_options
+
+  @options.setter
+  def options(self, options: QuantizationDebugOptions) -> None:
+    self._debug_options = options
+    if not self.converter or not self.calibrated_model:
+      return
+    self._init_from_converter(
+        self._debug_options,
+        self.converter,
+        self.calibrated_model,
+        float_model=self.float_model)
+    self._initialize_stats()
+
+  def _initialize_stats(self):
+    """Helper function initializes stats."""
+    # TODO(b/177749613) : Fix the dependency on tf.lite._get_ops_details()
+    # Following code is needed to get op's name from the output tensor index,
+    # since NumericVerify op only provides its quantized input tensor index.
+    self._defining_op = dict()
+    for op_info in self._quant_interpreter._get_ops_details():  # pylint: disable=protected-access
+      self._defining_op.update(
+          {tensor_idx: op_info['index'] for tensor_idx in op_info['outputs']})
+
+    self._numeric_verify_tensor_details = None
+    self._numeric_verify_op_details = None
+    if not self._get_numeric_verify_tensor_details():
+      raise ValueError('Please check if the quantized model is in debug mode')
+
+    self._layer_debug_metrics = _DEFAULT_LAYER_DEBUG_METRICS.copy()
+    if self._debug_options.layer_debug_metrics:
+      self._layer_debug_metrics.update(self._debug_options.layer_debug_metrics)
+
+    self.layer_statistics = None
+    self.model_statistics = None
+
+    self._metrics = metrics_stub.TFLiteMetrics()
+    self._metrics.increase_counter_debugger_creation()
+
+  def _get_quantized_model(self, is_debug: bool) -> bytes:
+    if not self.converter:
+      raise ValueError('No converter found, use this function with the '
+                       'converter option in the constructor.')
+
+    return convert.mlir_quantize(
+        self.calibrated_model,
+        disable_per_channel=self.converter._experimental_disable_per_channel,  # pylint: disable=protected-access
+        fully_quantize=self._debug_options.fully_quantize,
+        enable_numeric_verify=is_debug,
+        denylisted_ops=self._debug_options.denylisted_ops,
+        denylisted_nodes=self._debug_options.denylisted_nodes)
+
+  def get_nondebug_quantized_model(self) -> bytes:
+    """Returns a non-instrumented quantized model.
+
+    Convert the quantized model with the initialized converter and
+    return bytes for nondebug model. The model will not be instrumented with
+    numeric verification operations.
+
+    Returns:
+      Model bytes corresponding to the model.
+    Raises:
+      ValueError: if converter is not passed to the debugger.
+    """
+    return self._get_quantized_model(is_debug=False)
+
+  def get_debug_quantized_model(self) -> bytes:
+    """Returns an instrumented quantized model.
+
+    Convert the quantized model with the initialized converter and
+    return bytes for model. The model will be instrumented with numeric
+    verification operations and should only be used for debugging.
+
+    Returns:
+      Model bytes corresponding to the model.
+    Raises:
+      ValueError: if converter is not passed to the debugger.
+    """
+    return self._get_quantized_model(is_debug=True)
+
+  def _init_from_converter(self,
+                           options: QuantizationDebugOptions,
+                           converter: TFLiteConverter,
+                           calibrated_model: Optional[bytes] = None,
+                           float_model: Optional[bytes] = None) -> None:
+    """Convert the model and apply options.
+
+    Converts the quantized model and initializes a quantized model interpreter
+    with the quantized model. Returns a float model interpreter if float model
+    is provided.
+
+    Args:
+      options: a QuantizationDebugOptions object.
+      converter: an initialized tf.lite.TFLiteConverter.
+      calibrated_model: Calibrated model bytes.
+      float_model: Float model bytes.
+    """
+    self.quant_model = convert.mlir_quantize(
+        calibrated_model,
+        disable_per_channel=converter._experimental_disable_per_channel,  # pylint: disable=protected-access
+        fully_quantize=options.fully_quantize,
+        enable_numeric_verify=True,
+        denylisted_ops=options.denylisted_ops,
+        denylisted_nodes=options.denylisted_nodes)
+    self._quant_interpreter = _interpreter.Interpreter(
+        model_content=self.quant_model)
+    self._float_interpreter = None
+    if float_model is not None:
+      self._float_interpreter = _interpreter.Interpreter(
+          model_content=float_model)
+
+  def _set_converter_options_for_float(
+      self, converter: TFLiteConverter) -> TFLiteConverter:
+    """Verify converter options and set required experimental options."""
+    if converter.optimizations:
+      converter.optimizations = []
+    return converter
+
+  def _set_converter_options_for_calibration(
+      self, converter: TFLiteConverter) -> TFLiteConverter:
+    """Verify converter options and set required experimental options."""
+    if not converter.optimizations:
+      raise ValueError(
+          'converter object must set optimizations to lite.Optimize.DEFAULT')
+    if not converter.representative_dataset:
+      raise ValueError('converter object must set representative_dataset')
+
+    converter.experimental_mlir_quantizer = True
+    converter._experimental_calibrate_only = True  # pylint: disable=protected-access
+    return converter
+
+  def run(self) -> None:
+    """Runs models and gets metrics."""
+    self.layer_statistics = self._collect_layer_statistics()
+    if self._debug_options.model_debug_metrics:
+      self.model_statistics = self._collect_model_statistics()
+
+  def _collect_layer_statistics(self) -> Dict[str, Dict[str, float]]:
+    """Collects layer statistics by applying layer debug metrics.
+
+    For all data from the given RepresentativeDataset, collect statistics per
+    example by getting the NumericVerify op results in _quant_interpreter
+    and calculating layer debug metrics on the results.
+
+    Returns:
+      aggregated per-layer statistics of NumericVerify results.
+      {layer_name: {metric_name: metric}}
+    """
+    layer_statistics = collections.defaultdict(
+        lambda: collections.defaultdict(list))
+
+    initialize = True
+    for tensor_data in self._data_gen():
+      self._set_input_tensors(self._quant_interpreter, tensor_data, initialize)
+      initialize = False
+
+      # Run the model.
+      self._quant_interpreter.invoke()
+
+      # Collect the statistics of this invoke result.
+      for tensor_detail in self._get_numeric_verify_tensor_details():
+        tensor_name = tensor_detail['name']  # pytype: disable=unsupported-operands  # dynamic-method-lookup
+        diffs = self._quant_interpreter.get_tensor(tensor_detail['index'])  # pytype: disable=unsupported-operands  # dynamic-method-lookup
+        for metric_name, metric_fn in self._layer_debug_metrics.items():
+          layer_statistics[tensor_name][metric_name].append(metric_fn(diffs))
+
+      if self._debug_options.layer_direct_compare_metrics is not None:
+        for tensor_detail in self._get_numeric_verify_tensor_details():
+          tensor_name = tensor_detail['name']  # pytype: disable=unsupported-operands  # dynamic-method-lookup
+          op_idx = self._defining_op[tensor_detail['index']]  # pytype: disable=unsupported-operands  # dynamic-method-lookup
+          op_detail = self._quant_interpreter._get_op_details(op_idx)  # pylint: disable=protected-access
+          q_idx, f_idx = op_detail['inputs']
+          quant_input_detail = self._quant_interpreter._get_tensor_details(  # pylint: disable=protected-access
+              q_idx, subgraph_index=0)
+          for (metric_name, metric_fn
+              ) in self._debug_options.layer_direct_compare_metrics.items():
+            layer_statistics[tensor_name][metric_name].append(
+                metric_fn(
+                    self._quant_interpreter.get_tensor(f_idx),
+                    self._quant_interpreter.get_tensor(q_idx),
+                    quant_input_detail['quantization_parameters']['scales'][0],
+                    quant_input_detail['quantization_parameters']['zero_points']
+                    [0]))
+
+    # Calculate final aggregated metrics for each layer.
+    for metrics in layer_statistics.values():
+      for metric_name in metrics:
+        metrics[metric_name] = np.nanmean(metrics[metric_name])
+
+    return layer_statistics
+
+  def _collect_model_statistics(self) -> Dict[str, float]:
+    """Collects model output metrics.
+
+    For all data from the given RepresentativeDataset, collect all model output
+    results from float model & quantized debug model, and calculate metrics
+    by using model output functions. As a result, self.model_results is filled,
+
+    where self.model_results[model_output_function_name] = `aggregated model
+    output function value` (a scalar).
+
+    Returns:
+      aggregated per-model output discrepancy metrics.
+      {metric_name: aggregated_metric}
+    """
+
+    model_statistics = collections.defaultdict(list)
+
+    initialize = True
+    for tensor_data in self._data_gen():
+      # Run quantized debug model and collect output results.
+      self._set_input_tensors(self._quant_interpreter, tensor_data, initialize)
+      self._quant_interpreter.invoke()
+      quant_tensor_data = self._get_output_tensors(self._quant_interpreter)
+
+      # Run float model if it's initialized.
+      float_tensor_data = []
+      if self._float_interpreter:
+        self._set_input_tensors(
+            self._float_interpreter, tensor_data, initialize)
+        self._float_interpreter.invoke()
+        float_tensor_data = self._get_output_tensors(self._float_interpreter)
+
+      initialize = False
+
+      # Calculate the metrics.
+      for (metric_name,
+           metric_fn) in self._debug_options.model_debug_metrics.items():
+        model_statistics[metric_name].append(
+            metric_fn(float_tensor_data, quant_tensor_data))
+
+    # Calculate final aggregated metrics for each outputs.
+    return {
+        metric_name: np.mean(metric)
+        for metric_name, metric in model_statistics.items()
+    }
+
+  def _set_input_tensors(self, interpreter: _interpreter.Interpreter,
+                         tensor_data: Sequence[np.ndarray],
+                         initialize: bool) -> None:
+    """Sets input tensors into TFLite model Interpreter.
+
+    Args:
+      interpreter: a tf.lite.Interpreter object with allocated tensors.
+      tensor_data: a list of Numpy array data.
+      initialize: set to true when input is first set for the interpreter, to
+        set input shapes and allocate tensors.
+
+    Raises:
+      ValueError: when inputs can't be set, or size of provided inputs does not
+      match size of model inputs.
+    """
+    input_details = interpreter.get_input_details()
+    if len(input_details) != len(tensor_data):
+      raise ValueError(
+          'Number of inputs provided ({}) does not match number of inputs to '
+          'the model ({})'.format(len(tensor_data), len(input_details)))
+
+    if initialize:
+      for input_detail, tensor in zip(input_details, tensor_data):
+        interpreter.resize_tensor_input(input_detail['index'], tensor.shape)
+      interpreter.allocate_tensors()
+
+    for input_detail, tensor in zip(input_details, tensor_data):
+      if tensor.dtype == np.float32 and input_detail['dtype'] == np.int8:
+        quant_params = _get_quant_params(input_detail)
+        if quant_params:
+          scale, zero_point = quant_params
+          tensor = np.round((tensor / scale) + zero_point).astype(np.int8)
+      interpreter.set_tensor(input_detail['index'], tensor)
+
+  def _get_output_tensors(
+      self, interpreter: _interpreter.Interpreter) -> List[np.ndarray]:
+    """Returns output tensors of given TFLite model Interpreter.
+
+    Args:
+      interpreter: a tf.lite.Interpreter object with allocated tensors.
+
+    Returns:
+      a list of numpy arrays representing output tensor results.
+    """
+
+    outputs = []
+    for output_detail in interpreter.get_output_details():
+      tensor = interpreter.get_tensor(output_detail['index'])
+      if output_detail['dtype'] == np.int8:
+        quant_params = _get_quant_params(output_detail)
+        if quant_params:
+          scale, zero_point = quant_params
+          tensor = ((tensor.astype(np.float32) - zero_point) * scale).astype(
+              np.float32)
+      outputs.append(tensor)
+
+    return outputs
+
+  def _get_numeric_verify_tensor_details(self) -> List[str]:
+    """Returns all names of all tensors from NumericVerify op."""
+    # pylint: disable=protected-access
+    if not self._numeric_verify_tensor_details:
+      self._numeric_verify_tensor_details = []
+      self._numeric_verify_op_details = {}
+      for op_info in self._quant_interpreter._get_ops_details():
+        if op_info['op_name'] == _NUMERIC_VERIFY_OP_NAME:
+          self._numeric_verify_tensor_details.append(
+              self._quant_interpreter._get_tensor_details(
+                  op_info['outputs'][0], subgraph_index=0))
+          tensor_name = self._numeric_verify_tensor_details[-1]['name']
+          self._numeric_verify_op_details[tensor_name] = op_info
+    # pylint: enable=protected-access
+    return self._numeric_verify_tensor_details
+
+  def _get_operand_name_and_index(self,
+                                  numeric_verify_name: str) -> Tuple[str, int]:
+    """Gets the index and name of NumericVerify Op's quantized input tensor.
+
+    Args:
+      numeric_verify_name: name of the NumericVerify op's output tensor. It has
+        format of `NumericVerify/{quantized_tensor_name}:{quantized_tensor_idx}`
+
+    Returns:
+      Tuple of (tensor_name, tensor_idx) for quantized op's output tensor.
+    """
+    tensor_name, tensor_idx = numeric_verify_name.rsplit(':', 1)
+    float_tensor_name = tensor_name[len(_NUMERIC_VERIFY_OP_NAME) + 1:]
+    if re.match(r'\d', float_tensor_name[-1]):
+      float_tensor_name = float_tensor_name[:-1]
+
+    return (float_tensor_name, int(tensor_idx))
+
+  def layer_statistics_dump(self, file: IO[str]) -> None:
+    """Dumps layer statistics into file, in csv format.
+
+    Args:
+      file: file, or file-like object to write.
+    """
+    # order of `fields` is the order of fields in csv.
+    fields = ['op_name', 'tensor_idx'] + list(self._layer_debug_metrics.keys())
+    if self._debug_options.layer_direct_compare_metrics is not None:
+      fields += list(self._debug_options.layer_direct_compare_metrics.keys())
+    fields += ['scale', 'zero_point', 'tensor_name']
+    writer = csv.DictWriter(file, fields)
+    writer.writeheader()
+    if self.layer_statistics:
+      for name, metrics in self.layer_statistics.items():
+        data = metrics.copy()
+        (data['tensor_name'], _) = self._get_operand_name_and_index(name)
+        data['tensor_idx'] = self._numeric_verify_op_details[name]['inputs'][0]
+        data['op_name'] = self._quant_interpreter._get_op_details(  # pylint: disable=protected-access
+            self._defining_op[data['tensor_idx']])['op_name']
+        details = self._quant_interpreter._get_tensor_details(  # pylint: disable=protected-access
+            data['tensor_idx'], subgraph_index=0)
+        data['scale'], data['zero_point'] = (
+            details['quantization_parameters']['scales'][0],
+            details['quantization_parameters']['zero_points'][0])
+        writer.writerow(data)

tensorflow/lite/tools/optimize/debugging/python/debugger_test.py

@@ -0,0 +1,437 @@
+# Copyright 2021 The TensorFlow Authors. All Rights Reserved.
+#
+# 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 QuantizationDebugger."""
+
+import csv
+import io
+import re
+from unittest import mock
+
+from absl.testing import parameterized
+import numpy as np
+import tensorflow as tf
+
+from tensorflow.lite.python import convert
+from tensorflow.lite.python import lite
+from tensorflow.lite.python.metrics import metrics
+from tensorflow.lite.tools.optimize.debugging.python import debugger
+from tensorflow.python.framework import test_util
+from tensorflow.python.platform import test
+from tensorflow.python.trackable import autotrackable
+
+
+def _get_model():
+  """Returns somple model with Conv2D and representative dataset gen."""
+  root = autotrackable.AutoTrackable()
+  kernel_in = np.array([-2, -1, 1, 2], dtype=np.float32).reshape((2, 2, 1, 1))
+
+  @tf.function(
+      input_signature=[tf.TensorSpec(shape=[1, 3, 3, 1], dtype=tf.float32)])
+  def func(inp):
+    kernel = tf.constant(kernel_in, dtype=tf.float32)
+    conv = tf.nn.conv2d(inp, kernel, strides=1, padding='SAME')
+    output = tf.nn.relu(conv, name='output')
+    return output
+
+  root.f = func
+  to_save = root.f.get_concrete_function()
+  return (root, to_save)
+
+
+def _calibration_gen():
+  for i in range(5):
+    yield [np.arange(9).reshape((1, 3, 3, 1)).astype(np.float32) * i]
+
+
+def _convert_model(model, func):
+  """Converts TF model to TFLite float model."""
+  converter = lite.TFLiteConverterV2.from_concrete_functions([func], model)
+  # TODO(b/191205988): Explicitly disable saved model lowering in conversion.
+  converter.experimental_lower_to_saved_model = False
+  return converter.convert()
+
+
+def _quantize_converter(model, func, calibration_gen, debug=True):
+  """Returns a converter appropriate for the function and debug configs."""
+  converter = lite.TFLiteConverterV2.from_concrete_functions([func], model)
+  converter.target_spec.supported_ops = [lite.OpsSet.TFLITE_BUILTINS_INT8]
+  converter.representative_dataset = calibration_gen
+
+  # TODO(b/191205988): Explicitly disable saved model lowering in conversion.
+  converter.experimental_lower_to_saved_model = False
+
+  # Create a TFLite model with new quantizer and numeric verify ops.
+  converter.optimizations = [lite.Optimize.DEFAULT]
+  converter.experimental_new_quantizer = True
+  if debug:
+    converter._experimental_calibrate_only = True
+  return converter
+
+
+def _quantize_model(model,
+                    func,
+                    calibration_gen,
+                    quantized_io=False,
+                    debug=True):
+  """Quantizes model, in debug or normal mode."""
+  converter = _quantize_converter(model, func, calibration_gen, debug)
+  if debug:
+    calibrated = converter.convert()
+    return convert.mlir_quantize(
+        calibrated, enable_numeric_verify=True, fully_quantize=quantized_io)
+  else:
+    return converter.convert()
+
+
+def _dummy_fn(*unused_args):
+  return 0.0
+
+
+class QuantizationDebugOptionsTest(test_util.TensorFlowTestCase,
+                                   parameterized.TestCase):
+
+  @test_util.run_v2_only
+  def test_init_duplicate_keys_raises_ValueError(self):
+    with self.assertRaises(ValueError):
+      debugger.QuantizationDebugOptions(
+          layer_debug_metrics={
+              'a': _dummy_fn,
+              'b': _dummy_fn
+          },
+          model_debug_metrics={
+              'c': _dummy_fn,
+              'd': _dummy_fn
+          },
+          layer_direct_compare_metrics={
+              'a': _dummy_fn,
+              'e': _dummy_fn
+          })
+
+    with self.assertRaises(ValueError):
+      debugger.QuantizationDebugOptions(
+          layer_debug_metrics={
+              'a': _dummy_fn,
+              'b': _dummy_fn
+          },
+          layer_direct_compare_metrics={
+              'a': _dummy_fn,
+              'e': _dummy_fn
+          })
+
+
+class QuantizationDebuggerTest(test_util.TensorFlowTestCase,
+                               parameterized.TestCase):
+
+  @classmethod
+  def setUpClass(cls):
+    super().setUpClass()
+    cls.tf_model_root, cls.tf_model = _get_model()
+    cls.float_model = _convert_model(cls.tf_model_root, cls.tf_model)
+    cls.debug_model_float = _quantize_model(
+        cls.tf_model_root, cls.tf_model, _calibration_gen, quantized_io=False)
+    cls.debug_model_int8 = _quantize_model(
+        cls.tf_model_root, cls.tf_model, _calibration_gen, quantized_io=True)
+
+  @parameterized.named_parameters(
+      ('float_io', False, False),
+      ('quantized_io', True, False),
+      ('float_io_from_converter', False, True),
+      ('quantized_io_from_converter', True, True),
+  )
+  @test_util.run_v2_only
+  def test_layer_metrics(self, quantized_io, from_converter):
+    options = debugger.QuantizationDebugOptions(
+        layer_debug_metrics={'l1_norm': lambda diffs: np.mean(np.abs(diffs))})
+    if not from_converter:
+      if quantized_io:
+        debug_model = QuantizationDebuggerTest.debug_model_int8
+      else:
+        debug_model = QuantizationDebuggerTest.debug_model_float
+      quant_debugger = debugger.QuantizationDebugger(
+          quant_debug_model_content=debug_model,
+          debug_dataset=_calibration_gen,
+          debug_options=options)
+    else:
+      options.fully_quantize = quantized_io
+      quant_debugger = debugger.QuantizationDebugger(
+          converter=_quantize_converter(self.tf_model_root, self.tf_model,
+                                        _calibration_gen),
+          debug_dataset=_calibration_gen,
+          debug_options=options)
+
+    quant_debugger.run()
+
+    expected_quant_io_metrics = {
+        'num_elements': 9,
+        'stddev': 0.03850026,
+        'mean_error': 0.01673192,
+        'max_abs_error': 0.10039272,
+        'mean_squared_error': 0.0027558778,
+        'l1_norm': 0.023704167,
+    }
+    expected_float_io_metrics = {
+        'num_elements': 9,
+        'stddev': 0.050998904,
+        'mean_error': 0.007843441,
+        'max_abs_error': 0.105881885,
+        'mean_squared_error': 0.004357292,
+        'l1_norm': 0.035729896,
+    }
+    expected_metrics = (
+        expected_quant_io_metrics
+        if quantized_io else expected_float_io_metrics)
+    self.assertLen(quant_debugger.layer_statistics, 1)
+    actual_metrics = next(iter(quant_debugger.layer_statistics.values()))
+
+    self.assertCountEqual(expected_metrics.keys(), actual_metrics.keys())
+    for key, value in expected_metrics.items():
+      self.assertAlmostEqual(value, actual_metrics[key], places=5)
+
+    buffer = io.StringIO()
+    quant_debugger.layer_statistics_dump(buffer)
+    reader = csv.DictReader(buffer.getvalue().split())
+    actual_values = next(iter(reader))
+
+    expected_values = expected_metrics.copy()
+    expected_values.update({
+        'op_name': 'CONV_2D',
+        'tensor_idx': 7,
+        'scale': 0.15686275,
+        'zero_point': -128,
+        'tensor_name': r'Identity[1-9]?$'
+    })
+    for key, value in expected_values.items():
+      if isinstance(value, str):
+        self.assertIsNotNone(
+            re.match(value, actual_values[key]),
+            'String is different from expected string. Please fix test code if'
+            " it's being affected by graph manipulation changes.")
+      elif isinstance(value, list):
+        self.assertAlmostEqual(
+            value[0], float(actual_values[key][1:-1]), places=5)
+      else:
+        self.assertAlmostEqual(value, float(actual_values[key]), places=5)
+
+  @parameterized.named_parameters(
+      ('float_io', False),
+      ('quantized_io', True),
+  )
+  @test_util.run_v2_only
+  def test_model_metrics(self, quantized_io):
+    if quantized_io:
+      debug_model = QuantizationDebuggerTest.debug_model_int8
+    else:
+      debug_model = QuantizationDebuggerTest.debug_model_float
+    options = debugger.QuantizationDebugOptions(
+        model_debug_metrics={'stdev': lambda x, y: np.std(x[0] - y[0])})
+    quant_debugger = debugger.QuantizationDebugger(
+        quant_debug_model_content=debug_model,
+        float_model_content=QuantizationDebuggerTest.float_model,
+        debug_dataset=_calibration_gen,
+        debug_options=options)
+    quant_debugger.run()
+
+    expected_metrics = {'stdev': 0.050998904}
+    actual_metrics = quant_debugger.model_statistics
+
+    self.assertCountEqual(expected_metrics.keys(), actual_metrics.keys())
+    for key, value in expected_metrics.items():
+      self.assertAlmostEqual(value, actual_metrics[key], places=5)
+
+  @parameterized.named_parameters(
+      ('float_io', False),
+      ('quantized_io', True),
+  )
+  @test_util.run_v2_only
+  def test_layer_direct_compare_metrics(self, quantized_io):
+
+    def _corr(float_values, quant_values, scale, zero_point):
+      dequant_values = (quant_values.astype(np.int32) - zero_point) * scale
+      return np.corrcoef(float_values.flatten(), dequant_values.flatten())[0, 1]
+
+    if quantized_io:
+      debug_model = QuantizationDebuggerTest.debug_model_int8
+    else:
+      debug_model = QuantizationDebuggerTest.debug_model_float
+
+    options = debugger.QuantizationDebugOptions(
+        layer_direct_compare_metrics={'corr': _corr})
+    quant_debugger = debugger.QuantizationDebugger(
+        quant_debug_model_content=debug_model,
+        debug_dataset=_calibration_gen,
+        debug_options=options)
+    quant_debugger.run()
+
+    expected_metrics = {
+        'corr': 0.99999,
+    }
+    self.assertLen(quant_debugger.layer_statistics, 1)
+    actual_metrics = next(iter(quant_debugger.layer_statistics.values()))
+
+    for key, value in expected_metrics.items():
+      self.assertAlmostEqual(value, actual_metrics[key], places=4)
+
+  @test_util.run_v2_only
+  def test_wrong_input_raises_ValueError(self):
+
+    def wrong_calibration_gen():
+      for _ in range(5):
+        yield [
+            np.ones((1, 3, 3, 1), dtype=np.float32),
+            np.ones((1, 3, 3, 1), dtype=np.float32)
+        ]
+
+    quant_debugger = debugger.QuantizationDebugger(
+        quant_debug_model_content=QuantizationDebuggerTest.debug_model_float,
+        debug_dataset=wrong_calibration_gen)
+    with self.assertRaisesRegex(
+        ValueError, r'inputs provided \(2\).+inputs to the model \(1\)'):
+      quant_debugger.run()
+
+  @test_util.run_v2_only
+  def test_non_debug_model_raises_ValueError(self):
+    normal_quant_model = _quantize_model(
+        QuantizationDebuggerTest.tf_model_root,
+        QuantizationDebuggerTest.tf_model,
+        _calibration_gen,
+        debug=False)
+
+    with self.assertRaisesRegex(
+        ValueError, 'Please check if the quantized model is in debug mode'):
+      debugger.QuantizationDebugger(
+          quant_debug_model_content=normal_quant_model,
+          debug_dataset=_calibration_gen)
+
+  @parameterized.named_parameters(
+      ('empty quantization parameter', {
+          'quantization_parameters': {}
+      }, None),
+      ('empty scales/zero points', {
+          'quantization_parameters': {
+              'scales': [],
+              'zero_points': []
+          }
+      }, None),
+      ('invalid scales/zero points', {
+          'quantization_parameters': {
+              'scales': [1.0],
+              'zero_points': []
+          }
+      }, None),
+      ('correct case', {
+          'quantization_parameters': {
+              'scales': [0.5, 1.0],
+              'zero_points': [42, 7]
+          }
+      }, (0.5, 42)),
+  )
+  def test_get_quant_params(self, tensor_detail, expected_value):
+    self.assertEqual(debugger._get_quant_params(tensor_detail), expected_value)
+
+  @parameterized.named_parameters(
+      ('float_io', False),
+      ('quantized_io', True),
+  )
+  @test_util.run_v2_only
+  def test_denylisted_ops_from_option_setter(self, quantized_io):
+    options = debugger.QuantizationDebugOptions(
+        layer_debug_metrics={'l1_norm': lambda diffs: np.mean(np.abs(diffs))},
+        fully_quantize=quantized_io)
+    quant_debugger = debugger.QuantizationDebugger(
+        converter=_quantize_converter(self.tf_model_root, self.tf_model,
+                                      _calibration_gen),
+        debug_dataset=_calibration_gen,
+        debug_options=options)
+
+    options.denylisted_ops = ['CONV_2D']
+    # TODO(b/195084873): The exception is expected to check whether selective
+    # quantization was done properly, since after the selective quantization
+    # the model will have no quantized layers thus have no NumericVerify ops,
+    # resulted in this exception. Marked with a bug to fix this in more
+    # straightforward way.
+    with self.assertRaisesRegex(
+        ValueError, 'Please check if the quantized model is in debug mode'):
+      quant_debugger.options = options
+
+  @parameterized.named_parameters(
+      ('float_io', False),
+      ('quantized_io', True),
+  )
+  @test_util.run_v2_only
+  def test_denylisted_ops_from_option_constructor(self, quantized_io):
+    options = debugger.QuantizationDebugOptions(
+        layer_debug_metrics={'l1_norm': lambda diffs: np.mean(np.abs(diffs))},
+        fully_quantize=quantized_io,
+        denylisted_ops=['CONV_2D'])
+    # TODO(b/195084873): Count the number of NumericVerify op.
+    with self.assertRaisesRegex(
+        ValueError, 'Please check if the quantized model is in debug mode'):
+      _ = debugger.QuantizationDebugger(
+          converter=_quantize_converter(self.tf_model_root, self.tf_model,
+                                        _calibration_gen),
+          debug_dataset=_calibration_gen,
+          debug_options=options)
+
+  @parameterized.named_parameters(
+      ('float_io', False),
+      ('quantized_io', True),
+  )
+  @test_util.run_v2_only
+  def test_denylisted_nodes_from_option_setter(self, quantized_io):
+    options = debugger.QuantizationDebugOptions(
+        layer_debug_metrics={'l1_norm': lambda diffs: np.mean(np.abs(diffs))},
+        fully_quantize=quantized_io)
+    quant_debugger = debugger.QuantizationDebugger(
+        converter=_quantize_converter(self.tf_model_root, self.tf_model,
+                                      _calibration_gen),
+        debug_dataset=_calibration_gen,
+        debug_options=options)
+
+    options.denylisted_nodes = ['Identity']
+    # TODO(b/195084873): Count the number of NumericVerify op.
+    with self.assertRaisesRegex(
+        ValueError, 'Please check if the quantized model is in debug mode'):
+      quant_debugger.options = options
+
+  @parameterized.named_parameters(
+      ('float_io', False),
+      ('quantized_io', True),
+  )
+  @test_util.run_v2_only
+  def test_denylisted_nodes_from_option_constructor(self, quantized_io):
+    options = debugger.QuantizationDebugOptions(
+        layer_debug_metrics={'l1_norm': lambda diffs: np.mean(np.abs(diffs))},
+        fully_quantize=quantized_io,
+        denylisted_nodes=['Identity'])
+    # TODO(b/195084873): Count the number of NumericVerify op.
+    with self.assertRaisesRegex(
+        ValueError, 'Please check if the quantized model is in debug mode'):
+      _ = debugger.QuantizationDebugger(
+          converter=_quantize_converter(self.tf_model_root, self.tf_model,
+                                        _calibration_gen),
+          debug_dataset=_calibration_gen,
+          debug_options=options)
+
+  @mock.patch.object(metrics.TFLiteMetrics,
+                     'increase_counter_debugger_creation')
+  def test_creation_counter(self, increase_call):
+    debug_model = QuantizationDebuggerTest.debug_model_float
+    debugger.QuantizationDebugger(
+        quant_debug_model_content=debug_model, debug_dataset=_calibration_gen)
+    increase_call.assert_called_once()
+
+
+if __name__ == '__main__':
+  test.main()