Tom/tflite dequantize (#1322)

* Implement const dequantize pushing for per-axis dequantization

Signed-off-by: Tom Wildenhain <[email protected]>

* Add --dequantize flag

Signed-off-by: Tom Wildenhain <[email protected]>

* pylint fixes

Signed-off-by: Tom Wildenhain <[email protected]>

* Update readme

Signed-off-by: Tom Wildenhain <[email protected]>

Author: TomWildenhain-Microsoft

README.md

@@ -130,12 +130,14 @@ You find an end-to-end tutorial for ssd-mobilenet [here](tutorials/ConvertingSSD
 python -m tf2onnx.convert
     --saved-model SOURCE_SAVED_MODEL_PATH |
     --checkpoint SOURCE_CHECKPOINT_METAFILE_PATH |
+    --tflite SOURCE_TFLITE_PATH |
     --input | --graphdef SOURCE_GRAPHDEF_PB
     --output TARGET_ONNX_MODEL
     [--inputs GRAPH_INPUTS]
     [--outputs GRAPH_OUTPUS]
     [--inputs-as-nchw inputs_provided_as_nchw]
     [--opset OPSET]
+    [--dequantize]
     [--tag TAG]
     [--signature_def SIGNATURE_DEF]
     [--concrete_function CONCRETE_FUNCTION]
@@ -158,6 +160,12 @@ TensorFlow model as saved_model. We expect the path to the saved_model directory
 
 TensorFlow model as checkpoint. We expect the path to the .meta file.
 
+#### --tflite
+
+(This is experimental)
+
+Convert a tflite model by providing a path to the .tflite file. Inputs/outputs do not need to be specified.
+
 #### --input or --graphdef
 
 TensorFlow model as graphdef file.
@@ -182,6 +190,12 @@ ONNX requires default values for graph inputs to be constant, while Tensorflow's
 
 By default we use the opset 8 to generate the graph. By specifying ```--opset``` the user can override the default to generate a graph with the desired opset. For example ```--opset 5``` would create a onnx graph that uses only ops available in opset 5. Because older opsets have in most cases fewer ops, some models might not convert on a older opset.
 
+#### --dequantize
+
+(This is experimental, only supported for tflite)
+
+Produces a float32 model from a quantized tflite model. Detects ReLU and ReLU6 ops from quantization bounds.
+
 #### --tag
 
 Only valid with parameter `--saved_model`. Specifies the tag in the saved_model to be used. Typical value is 'serve'.

tests/ade20k.jpg

@@ -9,7 +9,7 @@
 from __future__ import unicode_literals
 
 # pylint: disable=broad-except,logging-not-lazy,unused-argument,unnecessary-lambda,import-outside-toplevel
-# pylint: disable=wrong-import-position
+# pylint: disable=wrong-import-position,too-many-nested-blocks
 
 import argparse
 import os
@@ -79,6 +79,16 @@ def get_car(shape):
     return get_img(shape, "car.JPEG", np.float32, should_scale=True)
 
 
+def get_ade20k(shape):
+    """Get truck image from ade20k segmentation dataset."""
+    return get_img(shape, "ade20k.jpg", np.float32, should_scale=True)
+
+
+def get_ade20k_uint8(shape):
+    """Get truck image from ade20k segmentation dataset."""
+    return get_img(shape, "ade20k.jpg", np.uint8, should_scale=False)
+
+
 def get_random(shape):
     """Get random input."""
     np.random.seed(42)
@@ -146,6 +156,8 @@ def get_sentence():
 _INPUT_FUNC_MAPPING = {
     "get_beach": get_beach,
     "get_car": get_car,
+    "get_ade20k": get_ade20k,
+    "get_ade20k_uint8": get_ade20k_uint8,
     "get_random": get_random,
     "get_random256": get_random256,
     "get_ramp": get_ramp,
@@ -171,7 +183,7 @@ class Test(object):
     target = []
 
     def __init__(self, url, local, input_func, input_names, output_names,
-                 disabled=False, rtol=0.01, atol=1e-6,
+                 disabled=False, rtol=0.01, atol=1e-6, ptol=0, dequantize=False,
                  check_only_shape=False, model_type="frozen", force_input_shape=False,
                  skip_tensorflow=False, opset_constraints=None, tf_min_version=None, tag=None,
                  skip_conversion=False, converted_model=None, signature_def=None, concrete_function=None,
@@ -190,6 +202,8 @@ def __init__(self, url, local, input_func, input_names, output_names,
         self.structured_outputs = structured_outputs  # Needed to determine output order for tf_function
         self.rtol = rtol
         self.atol = atol
+        self.ptol = ptol
+        self.dequantize = dequantize
         self.check_only_shape = check_only_shape
         self.perf = None
         self.tf_runtime = 0
@@ -292,7 +306,7 @@ def to_onnx(self, tf_graph, opset=None, extra_opset=None, shape_override=None, i
                                 extra_opset=extra_opset, target=Test.target, shape_override=shape_override,
                                 input_names=input_names, output_names=self.output_names,
                                 const_node_values=const_node_values, initialized_tables=initialized_tables,
-                                tflite_path=tflite_path)
+                                tflite_path=tflite_path, dequantize=self.dequantize)
 
     def run_caffe2(self, name, model_proto, inputs):
         """Run test again caffe2 backend."""
@@ -531,7 +545,11 @@ def run_tflite():
                             np.testing.assert_array_equal(tf_res.shape, onnx_res.shape)
                     else:
                         for tf_res, onnx_res in zip(tf_results, onnx_results):
-                            np.testing.assert_allclose(tf_res, onnx_res, rtol=self.rtol, atol=self.atol)
+                            good_cnt = np.count_nonzero(np.isclose(tf_res, onnx_res, rtol=self.rtol, atol=self.atol))
+                            bad_cnt = tf_res.size - good_cnt
+                            if bad_cnt > self.ptol / 100 * tf_res.size:
+                                # Prints a nice error message with stats
+                                np.testing.assert_allclose(tf_res, onnx_res, rtol=self.rtol, atol=self.atol)
                     logger.info("Results: OK")
                 return True
             except Exception:
@@ -658,10 +676,10 @@ def load_tests_from_yaml(path):
                 opset_constraints.append(c)
 
         kwargs = {}
-        for kw in ["rtol", "atol", "disabled", "check_only_shape", "model_type", "concrete_function",
+        for kw in ["rtol", "atol", "ptol", "disabled", "check_only_shape", "model_type", "concrete_function",
                    "skip_tensorflow", "force_input_shape", "tf_min_version", "tag", "skip_conversion",
                    "converted_model", "signature_def", "large_model", "structured_outputs", "run_tf_frozen",
-                   "use_custom_ops"]:
+                   "use_custom_ops", "dequantize"]:
             if settings.get(kw) is not None:
                 kwargs[kw] = settings[kw]
 

tests/run_pretrained_models.py

@@ -454,3 +454,30 @@ ssd_mobilenet_v2_300_float_tflite:
     - TFLite_Detection_PostProcess:1
     - TFLite_Detection_PostProcess:2
     - TFLite_Detection_PostProcess:3
+
+deeplabv3_mnv2_ade20k_float_tflite:
+  tf_min_version: 2.1
+  disabled: false
+  url: https://github.com/mlcommons/mobile_models/raw/main/v0_7/tflite/deeplabv3_mnv2_ade20k_float.tflite
+  model: "deeplabv3_mnv2_ade20k_float.tflite"
+  model_type: tflite
+  input_get: get_ade20k
+  ptol: 0.001
+  inputs:
+    "MobilenetV2/MobilenetV2/input": [1, 512, 512, 3]
+  outputs:
+    - ArgMax
+
+deeplabv3_mnv2_ade20k_uint8_tflite:
+  tf_min_version: 2.1
+  disabled: false
+  url: https://github.com/mlcommons/mobile_models/raw/main/v0_7/tflite/deeplabv3_mnv2_ade20k_uint8.tflite
+  model: "deeplabv3_mnv2_ade20k_uint8.tflite"
+  model_type: tflite
+  input_get: get_ade20k_uint8
+  ptol: 1.0
+  dequantize: true
+  inputs:
+    "MobilenetV2/MobilenetV2/input": [1, 512, 512, 3]
+  outputs:
+    - ArgMax

tests/run_pretrained_models.yaml

@@ -65,6 +65,8 @@ def get_args():
     parser.add_argument("--use_default", help="comma-separated list of names of PlaceholderWithDefault ops to "
                                               "change into Identity ops using their default value")
     parser.add_argument("--opset", type=int, default=None, help="opset version to use for onnx domain")
+    parser.add_argument("--dequantize", help="Remove quantization from model. Only supported for tflite currently.",
+                        action="store_true")
     parser.add_argument("--custom-ops", help="comma-separated map of custom ops to domains in format OpName:domain")
     parser.add_argument("--extra_opset", default=None,
                         help="extra opset with format like domain:version, e.g. com.microsoft:1")
@@ -104,6 +106,9 @@ def get_args():
         args.target = args.target.split(",")
     if args.signature_def:
         args.signature_def = [args.signature_def]
+    if args.dequantize:
+        if not args.tflite:
+            parser.error("dequantize flag is currently only supported for tflite")
     if args.extra_opset:
         tokens = args.extra_opset.split(':')
         if len(tokens) != 2:
@@ -202,7 +207,8 @@ def main():
                              use_default=args.use_default,
                              const_node_values=const_node_values,
                              initialized_tables=initialized_tables,
-                             tflite_path=tflite_path)
+                             tflite_path=tflite_path,
+                             dequantize=args.dequantize)
 
     onnx_graph = optimizer.optimize_graph(g)
 

tf2onnx/convert.py

@@ -7,6 +7,7 @@
 
 import logging
 import numpy as np
+from onnx.onnx_pb import TensorProto
 from tf2onnx.handler import tfl_op
 from tf2onnx import utils
 
@@ -87,31 +88,76 @@ def to_tf(cls, ctx, node, **kwargs):
 
 @tfl_op(["TFL_QUANTIZE"], onnx_op="QuantizeLinear")
 class TflQuantizeOp:
+    @classmethod
+    def version_1(cls, ctx, node, dequantize=False, **kwargs):
+        # We could just let the TFL_QUANTIZE fall through as an unconverted op, but they are added programmatically
+        # so that might be confusing.
+        raise ValueError("Opset 10 is required for quantization. Consider using the --dequantize flag or --opset 10.")
+
     @classmethod
     def version_10(cls, ctx, node, **kwargs):
         scale = node.get_attr_value('scale')
         zero_point = node.get_attr_value('zero_point')
         axis = node.get_attr_value('quantized_dimension')
         np_q_type = utils.map_onnx_to_numpy_type(ctx.get_dtype(node.output[0]))
         if len(scale) > 1 or len(zero_point) > 1:
+            utils.make_sure(ctx.opset >= 13, "Opset 13 is required for per-axis quantization for node %s", node.name)
             node.set_attr("axis", axis)
         scale_node = ctx.make_const(utils.make_name("scale"), np.array(scale[0], dtype=np.float32))
         zero_point_node = ctx.make_const(utils.make_name("zero_point"), np.array(zero_point[0], dtype=np_q_type))
         ctx.replace_inputs(node, [node.input[0], scale_node.output[0], zero_point_node.output[0]])
         del node.attr["scale"]
         del node.attr["zero_point"]
         del node.attr["quantized_dimension"]
+        if "min" in node.attr:
+            del node.attr["min"]
+        if "max" in node.attr:
+            del node.attr["max"]
 
 @tfl_op(["TFL_DEQUANTIZE"], onnx_op="DequantizeLinear")
 class TflDequantizeOp:
     @classmethod
-    def version_10(cls, ctx, node, **kwargs):
+    def version_1(cls, ctx, node, **kwargs):
+        scale = np.array(node.get_attr_value('scale'), dtype=np.float32)
+        zero_point = np.array(node.get_attr_value('zero_point'), dtype=np.float32)
+        axis = node.get_attr_value('quantized_dimension')
+        in_rank = ctx.get_rank(node.input[0])
+        def expand_tensor(t):
+            if t.shape == (1,):
+                return t[0]
+            utils.make_sure(in_rank is not None, "Cannot dequantize node %s with unknown input rank", node.name)
+            new_shape = [1] * in_rank
+            new_shape[axis] = t.shape[0]
+            return t.reshape(new_shape)
+        scale = expand_tensor(scale)
+        zero_point = expand_tensor(zero_point)
+        if node.inputs[0].is_const():
+            x_val = node.inputs[0].get_tensor_value(as_list=False).astype(np.float32)
+            new_val = (x_val - zero_point) * scale
+            dequant_const = ctx.make_const(utils.make_name(node.name), new_val)
+            ctx.replace_all_inputs(node.output[0], dequant_const.output[0])
+            ctx.remove_node(node.name)
+        else:
+            scale_const = ctx.make_const(utils.make_name(node.name + "_scale"), scale).output[0]
+            zero_point_const = ctx.make_const(utils.make_name(node.name + "_zero_point"), zero_point).output[0]
+            cast_node = ctx.make_node("Cast", [node.input[0]], attr={'to': TensorProto.FLOAT},
+                                      op_name_scope=node.name).output[0]
+            sub_node = ctx.make_node("Sub", [cast_node, zero_point_const], op_name_scope=node.name).output[0]
+            mul_node = ctx.make_node("Mul", [sub_node, scale_const], op_name_scope=node.name).output[0]
+            ctx.replace_all_inputs(node.output[0], mul_node)
+            ctx.remove_node(node.name)
+
+    @classmethod
+    def version_10(cls, ctx, node, dequantize=False, **kwargs):
+        if dequantize:
+            cls.version_1(ctx, node, dequantize=True, **kwargs)
+            return
         scale = node.get_attr_value('scale')
         zero_point = node.get_attr_value('zero_point')
         axis = node.get_attr_value('quantized_dimension')
         np_q_type = utils.map_onnx_to_numpy_type(ctx.get_dtype(node.input[0]))
         if len(scale) > 1 or len(zero_point) > 1:
-            utils.make_sure(ctx.opset >= 13, "Opset 13 is required for per-axis quantization")
+            utils.make_sure(ctx.opset >= 13, "Opset 13 is required for per-axis quantization for node %s", node.name)
             node.set_attr("axis", axis)
             scale_node = ctx.make_const(utils.make_name("scale"), np.array(scale, dtype=np.float32))
             zero_point_node = ctx.make_const(utils.make_name("zero_point"), np.array(zero_point, dtype=np_q_type))
@@ -122,6 +168,10 @@ def version_10(cls, ctx, node, **kwargs):
         del node.attr["scale"]
         del node.attr["zero_point"]
         del node.attr["quantized_dimension"]
+        if "min" in node.attr:
+            del node.attr["min"]
+        if "max" in node.attr:
+            del node.attr["max"]
 
 def dynamic_quantize_inputs(ctx, node):
     if ctx.opset < 11:

tf2onnx/tflite_handlers/tfl_math.py

@@ -3,7 +3,9 @@
 """tf2onnx.tflite_rewriters module"""
 
 from tf2onnx.tflite_rewriters.tfl_scan_output_rewriter import rewrite_tfl_scan_outputs
+from tf2onnx.tflite_rewriters.tfl_qdq_rewriter import rewrite_tfl_qdq
 
 __all__ = [
     "rewrite_tfl_scan_outputs",
+    "rewrite_tfl_qdq"
 ]

tf2onnx/tflite_rewriters/__init__.py

@@ -0,0 +1,57 @@
+# SPDX-License-Identifier: Apache-2.0
+
+
+"""
+tf2onnx.tflite_rewriters.tfl_qdq_rewriter - Remove qdq sequences to dequantize model
+"""
+from tf2onnx.graph_matcher import OpTypePattern, GraphMatcher
+
+
+# pylint: disable=missing-docstring
+
+def rewrite_tfl_qdq(g, ops):
+    pattern0 = \
+        OpTypePattern('TFL_DEQUANTIZE', name='dequant', inputs=[
+            OpTypePattern('TFL_QUANTIZE', name='quant'),
+        ])
+
+    matcher = GraphMatcher(pattern0, allow_reorder=False)
+    match_results = list(matcher.match_ops(ops))
+    if match_results:
+        for match in match_results:
+            dequant = match.get_op("dequant")
+            quant = match.get_op("quant")
+            inp_node = quant.inputs[0]
+            for k in ["scale", "quantized_dimension", "zero_point"]:
+                if dequant.get_attr_value(k) != quant.get_attr_value(k):
+                    continue
+            needed_relu = None
+            if all(k in quant.attr and len(quant.get_attr_value(k)) == 1 for k in ["min", "max"]):
+                min_val = quant.get_attr_value("min")[0]
+                max_val = quant.get_attr_value("max")[0]
+                if min_val == 0.0 and 5.999 <= max_val <= 6.0:
+                    needed_relu = "TFL_RELU6"
+                elif min_val == 0.0:
+                    # This may introduce unneeded relu ops but will be correct.
+                    # If the --dequantize feature is used a lot in the future we can optimize this.
+                    needed_relu = "TFL_RELU"
+                if inp_node.type == needed_relu:
+                    # If it's really obviously unneeded, we skip it.
+                    needed_relu = None
+                elif "TFL_" + inp_node.get_attr_value("fused_activation_function", b'').decode() == needed_relu:
+                    needed_relu = None
+
+            if needed_relu is not None:
+                relu_name = inp_node.name + "_relu"
+
+                relu6 = g.make_node(needed_relu, [quant.input[0]], op_name_scope=relu_name,
+                                    skip_conversion=False, shapes=quant.output_shapes, dtypes=quant.output_dtypes)
+                g.replace_all_inputs(dequant.output[0], relu6.output[0])
+            else:
+                g.replace_all_inputs(dequant.output[0], quant.input[0])
+
+            g.remove_node(dequant.name)
+            if len(g.find_output_consumers(quant.output[0])) == 0:
+                g.remove_node(quant.name)
+
+    return ops