onnx2tf 1.29.14__py3-none-any.whl → 1.29.16__py3-none-any.whl

onnx2tf/ops/QLinearConcat.py

@@ -10,6 +10,10 @@ from onnx2tf.utils.common_functions import (
     print_node_info,
     inverted_operation_enable_disable,
     make_tf_node_info,
+    pre_process_transpose,
+    replace_parameter,
+    shape_is_equal_ignore_order,
+    transpose_with_flexing_deterrence,
 )
 
 
@@ -35,8 +39,7 @@ def make_node(
     y_zero_point_list = [i for i in graph_node.inputs[1::3]]
     input_list = [i for i in graph_node.inputs[2::3]]
 
-    input_tensor_shape = input_list[0].shape
-    input_tensor_rank = len(input_tensor_shape)
+    input_tensor_rank = len(input_list[0].shape)
 
     before_op_output_shape_trans = True
     for graph_node_input in input_list:
@@ -46,6 +49,9 @@ def make_node(
             before_op_output_shape_trans and before_op_output_shape_trans_n
 
     got_values = []
+    nhwc_flags = []
+    same_input_shape_as_onnxs = []
+    input_is_dequantized_list = []
     got_y_scale_list = []
     got_y_zero_point_list = []
     for input, y_scale, y_zero_point  in zip(input_list, y_scale_list, y_zero_point_list):
@@ -55,8 +61,24 @@ def make_node(
         )
         if isinstance(const_or_var, gs.Variable):
             got_values.append(tf_layers_dict[const_or_var.name]['tf_node'])
+            nhwc_flags.append(
+                tf_layers_dict[const_or_var.name].get('nhwc', False)
+            )
+            same_input_shape_as_onnxs.append(
+                True if input.shape is not None and len(input.shape) > 0 \
+                    and input.shape == tf_layers_dict[const_or_var.name]['tf_node'].shape else False
+            )
+            input_is_dequantized_list.append(
+                tf_layers_dict[const_or_var.name].get('is_dequantized', False)
+            )
         else:
             got_values.append(const_or_var)
+            nhwc_flags.append(False)
+            same_input_shape_as_onnxs.append(
+                True if input.shape is not None and len(input.shape) > 0 \
+                    and input.shape == const_or_var.shape else False
+            )
+            input_is_dequantized_list.append(False)
 
         const_or_var = get_constant_or_variable(
             y_scale,
@@ -82,50 +104,247 @@ def make_node(
     dtype = graph_node_output.dtype
 
     axis = graph_node.attrs.get('axis', 0)
-    # NCHW->NHWC, NCDHW->NDHWC
-    axis = convert_axis(
-        axis=axis,
-        tensor_rank=len(shape) if shape is not None else input_tensor_rank,
-        before_op_output_shape_trans=before_op_output_shape_trans,
-    )
+
+    # Shape Unmatched Special Avoidance Workaround
+    if True in same_input_shape_as_onnxs and True in nhwc_flags:
+        before_op_output_shape_trans = True
+        new_values = []
+        for same_input_shape_as_onnx, nhwc_flag, value in zip(same_input_shape_as_onnxs, nhwc_flags, got_values):
+            if same_input_shape_as_onnx and not nhwc_flag:
+                if len(value.shape) == 3:
+                    new_values.append(
+                        transpose_with_flexing_deterrence(
+                            input_tensor=value,
+                            perm=[0, 2, 1],
+                            **kwargs,
+                        )
+                    )
+                elif len(value.shape) == 4:
+                    new_values.append(
+                        transpose_with_flexing_deterrence(
+                            input_tensor=value,
+                            perm=[0, 2, 3, 1],
+                            **kwargs,
+                        )
+                    )
+                elif len(value.shape) == 5:
+                    new_values.append(
+                        transpose_with_flexing_deterrence(
+                            input_tensor=value,
+                            perm=[0, 2, 3, 4, 1],
+                            **kwargs,
+                        )
+                    )
+                else:
+                    new_values.append(value)
+            else:
+                new_values.append(value)
+        got_values = new_values
 
     # Preserving Graph Structure (Dict)
+    nhwc_judge = True
+    for graph_node_input in input_list:
+        if isinstance(graph_node_input, gs.Variable) \
+            and tf_layers_dict.get(graph_node_input.name, {}).get('nhwc', False):
+            nhwc_judge = nhwc_judge and True
+        elif isinstance(graph_node_input, gs.Constant) \
+            and hasattr(graph_node_input, 'values') \
+            and isinstance(graph_node_input.values, np.ndarray):
+            nhwc_judge = nhwc_judge or False
+        else:
+            nhwc_judge = nhwc_judge and False
+
     tf_layers_dict[graph_node_output.name] = {
         'optype': graph_node.op,
         'shape': shape,
         'dtype': dtype,
+        'is_dequantized': True,
     }
+    if nhwc_judge:
+        tf_layers_dict[graph_node_output.name]['nhwc'] = True
 
     # Generation of TF OP
 
+    # NCHW->NHWC, NCDHW->NDHWC
+    axis = convert_axis(
+        axis=axis,
+        tensor_rank=len(shape) if shape is not None else input_tensor_rank,
+        before_op_output_shape_trans=before_op_output_shape_trans,
+    )
+
+    # Param replacement
+    before_axis = axis
+    axis = replace_parameter(
+        value_before_replacement=axis,
+        param_target='attributes',
+        param_name='axis',
+        **kwargs,
+    )
+
     # TensorFlow does not support Concat for scalar values, so convert to tensor
-    values = [
-        value if len(value.shape) > 0 else tf.reshape(value, [1]) for value in got_values
-    ]
+    values = []
+    for graph_node_input, value in zip(input_list, got_values):
+        value = pre_process_transpose(
+            value_before_transpose=value,
+            param_target='inputs',
+            param_name=graph_node_input.name,
+            **kwargs,
+        )
+        values.append(value if len(value.shape) > 0 else tf.reshape(value, [1]))
+
+    def _infer_concat_axis(values, output_shape):
+        if not values:
+            return None
+        ranks = []
+        shapes = []
+        for val in values:
+            if val.shape is None or val.shape == tf.TensorShape(None):
+                return None
+            shape_list = list(val.shape)
+            ranks.append(len(shape_list))
+            shapes.append(shape_list)
+        if len(set(ranks)) != 1:
+            return None
+        rank = ranks[0]
+        candidates = []
+        for ax in range(rank):
+            ok = True
+            for dim in range(rank):
+                if dim == ax:
+                    continue
+                base = shapes[0][dim]
+                for s in shapes[1:]:
+                    if base is None or s[dim] is None:
+                        continue
+                    if base != s[dim]:
+                        ok = False
+                        break
+                if not ok:
+                    break
+            if not ok:
+                continue
+            if output_shape is not None and len(output_shape) == rank:
+                out_dim = output_shape[ax]
+                if out_dim is not None:
+                    sum_dim = 0
+                    for s in shapes:
+                        if s[ax] is None:
+                            sum_dim = None
+                            break
+                        sum_dim += s[ax]
+                    if sum_dim is None or sum_dim != out_dim:
+                        continue
+            candidates.append(ax)
+        if len(candidates) == 1:
+            return candidates[0]
+        return None
+
+    inferred_axis = _infer_concat_axis(values, shape if shape is not None else None)
+    if inferred_axis is not None:
+        axis = inferred_axis
     # cast all inputs to float32
     casted_x_list = []
     casted_y_zero_point_list = []
-    for x, y_zero_point in zip(values, got_y_zero_point_list):
+    casted_y_scale_list = []
+    for x, y_scale, y_zero_point in zip(values, got_y_scale_list, got_y_zero_point_list):
         casted_x_list.append(tf.cast(x, tf.float32))
+        casted_y_scale_list.append(tf.cast(y_scale, tf.float32))
         casted_y_zero_point_list.append(tf.cast(y_zero_point, tf.float32))
     # dequantize x with y_scale, y_zero_point
     dequantized_x_list = []
-    for x, y_scale, y_zero_point in zip(casted_x_list, got_y_scale_list, casted_y_zero_point_list):
-        dequantized_value = tf.add(
-            x=tf.divide(
-                x=x,
+    for x, y_scale, y_zero_point, is_dequantized in zip(
+        casted_x_list,
+        casted_y_scale_list,
+        casted_y_zero_point_list,
+        input_is_dequantized_list,
+    ):
+        if is_dequantized:
+            dequantized_x_list.append(x)
+        else:
+            dequantized_value = tf.multiply(
+                x=tf.subtract(x, y_zero_point),
                 y=y_scale,
-            ),
-            y=y_zero_point,
-        )
-        dequantized_x_list.append(dequantized_value)
+            )
+            dequantized_x_list.append(dequantized_value)
 
-    tf_layers_dict[graph_node_output.name]['tf_node'] = \
-        tf.concat(
-            values=dequantized_x_list,
-            axis=axis,
-            name=graph_node.name,
-        )
+    try:
+        tf_layers_dict[graph_node_output.name]['tf_node'] = \
+            tf.concat(
+                values=dequantized_x_list,
+                axis=axis,
+                name=graph_node.name,
+            )
+    except:
+        try:
+            onnx_axis = int(graph_node.attrs.get('axis', 0))
+            tf_layers_dict[graph_node_output.name]['tf_node'] = \
+                tf.concat(
+                    values=dequantized_x_list,
+                    axis=onnx_axis,
+                    name=graph_node.name,
+                )
+            axis = onnx_axis
+        except:
+            value_rank = len(dequantized_x_list[0].shape)
+            succeed = False
+            for idx in reversed(range(value_rank)):
+                try:
+                    tf_layers_dict[graph_node_output.name]['tf_node'] = \
+                        tf.concat(
+                            values=dequantized_x_list,
+                            axis=idx,
+                            name=graph_node.name,
+                        )
+                    axis = idx
+                    succeed = True
+                    break
+                except:
+                    pass
+            if not succeed:
+                raise
+
+    output_tensor_shape = tf_layers_dict[graph_node_output.name]['tf_node'].shape
+    if output_tensor_shape != tf.TensorShape(None):
+        output_tensor_rank = len(output_tensor_shape)
+        if graph_node.outputs[0].shape is not None \
+            and axis != 0 \
+            and output_tensor_rank >= 2 \
+            and before_axis == axis:
+            if not shape_is_equal_ignore_order(list(graph_node.outputs[0].shape), list(output_tensor_shape)):
+                matched_axes = []
+                for dummy_axis in range(1, output_tensor_rank):
+                    try:
+                        dummy_concat_tensor = \
+                            tf.concat(
+                                values=dequantized_x_list,
+                                axis=dummy_axis,
+                                name=graph_node.name,
+                            )
+                        dummy_output_shape = dummy_concat_tensor.shape
+                        if shape_is_equal_ignore_order(list(graph_node.outputs[0].shape), list(dummy_output_shape)):
+                            matched_axes.append(dummy_axis)
+                    except:
+                        pass
+                if len(matched_axes) == 1:
+                    tf_layers_dict[graph_node_output.name]['tf_node'] = \
+                        tf.concat(
+                            values=dequantized_x_list,
+                            axis=matched_axes[0],
+                            name=graph_node.name,
+                        )
+                    axis = matched_axes[0]
+                elif not nhwc_judge:
+                    onnx_axis = int(graph_node.attrs.get('axis', 0))
+                    onnx_axis = output_tensor_rank - 1 if onnx_axis == -1 else onnx_axis
+                    if onnx_axis == output_tensor_rank - 1 \
+                        and onnx_axis in matched_axes:
+                        tf_layers_dict[graph_node_output.name]['tf_node'] = \
+                            tf.concat(
+                                values=dequantized_x_list,
+                                axis=onnx_axis,
+                                name=graph_node.name,
+                            )
+                        axis = onnx_axis
 
     # Generation of Debug Info
     tf_inputs = {f"input{idx}": dequantized_x for idx, dequantized_x in enumerate(dequantized_x_list)}

onnx2tf/ops/QLinearConv.py

@@ -1,4 +1,3 @@
-import sys
 import random
 random.seed(0)
 import numpy as np
@@ -25,45 +24,57 @@ def _dequantize_tensor(
     # Do computation in float32
     base = tf.cast(base, tf.float32)
     zero_point = tf.cast(zero_point, tf.float32)
+    scale = tf.cast(scale, tf.float32)
     return (base - zero_point) * scale
 
 
+def _reshape_per_output_channel(
+    *,
+    value,
+    weights,
+):
+    value_rank = len(value.shape)
+    weights_rank = len(weights.shape)
+    if value_rank == 0:
+        return value
+    if value_rank == 1 and weights_rank is not None:
+        shape = [1] * weights_rank
+        shape[-1] = -1
+        return tf.reshape(value, shape)
+    return value
+
+
 def _dequantize_weights(
     *,
     base,
     zero_point,
     scale,
-    is_bias=False,
-    scale_is_scalar=False,
 ):
     # Do computation in float32
     casted_base = tf.cast(base, tf.float32)
     casted_zero_point = tf.cast(zero_point, tf.float32)
-    spartial_shape_len = len(casted_base.shape) - 2
-    casted_zero_point_shape = casted_zero_point.shape[0]
-    if casted_zero_point_shape == base.shape[-2]:
-        reshaped_zero_point = tf.reshape(
-            tensor=casted_zero_point,
-            shape=[1 for _ in range(spartial_shape_len)] + [casted_zero_point_shape, 1],
-        )
-        if scale_is_scalar:
-            reshaped_scale = tf.reshape(
-                tensor=scale,
-                shape=[1 for _ in range(spartial_shape_len)] + [casted_zero_point_shape, 1],
-            )
-            tensor_list = [
-                (casted_base[..., i:i+1] - reshaped_zero_point) * reshaped_scale
-                for i in range(base.shape[-1])
-            ]
-            out_tensor = tf.concat(tensor_list, axis=-1)
-        else:
-            reshaped_scale = scale
-            out_tensor = (casted_base - reshaped_zero_point) * reshaped_scale
-        return tf.reshape(out_tensor, base.shape)
-    else:
-        reshaped_zero_point = casted_zero_point
-        reshaped_scale = scale
-        return (casted_base - reshaped_zero_point) * reshaped_scale
+    casted_scale = tf.cast(scale, tf.float32)
+    casted_zero_point = _reshape_per_output_channel(
+        value=casted_zero_point,
+        weights=casted_base,
+    )
+    casted_scale = _reshape_per_output_channel(
+        value=casted_scale,
+        weights=casted_base,
+    )
+    return (casted_base - casted_zero_point) * casted_scale
+
+
+def _get_qmin_qmax(dtype: tf.dtypes.DType):
+    if dtype == tf.uint8:
+        return 0.0, 255.0
+    if dtype == tf.int8:
+        return -128.0, 127.0
+    if dtype == tf.uint16:
+        return 0.0, 65535.0
+    if dtype == tf.int16:
+        return -32768.0, 32767.0
+    return None, None
 
 
 @print_node_info
@@ -139,6 +150,11 @@ def make_node(
 
     input_tensor = tf_layers_dict[graph_node_input_1.name]['tf_node'] \
         if isinstance(graph_node_input_1, gs.Variable) else graph_node_input_1
+    input_is_dequantized = False
+    input_nhwc = False
+    if isinstance(graph_node_input_1, gs.Variable):
+        input_is_dequantized = tf_layers_dict.get(graph_node_input_1.name, {}).get('is_dequantized', False)
+        input_nhwc = tf_layers_dict.get(graph_node_input_1.name, {}).get('nhwc', False)
     input_tensor_scale = tf_layers_dict[graph_node_input_2.name]['tf_node'] \
         if isinstance(graph_node_input_2, gs.Variable) else graph_node_input_2
     input_tensor_zero_point = tf_layers_dict[graph_node_input_3.name]['tf_node'] \
@@ -155,7 +171,7 @@ def make_node(
         if isinstance(graph_node_input_8, gs.Variable) else graph_node_input_8
     input_bias = tf_layers_dict[graph_node_input_9.name]['tf_node'] \
         if isinstance(graph_node_input_9, gs.Variable) else graph_node_input_9
-    output_dtype = input_tensor.dtype if input_tensor.dtype not in [tf.int8, tf.uint8] else tf.float32
+    output_quant_dtype = y_zero_point.dtype
 
     input_tensor_shape = input_tensor.shape
     input_tensor_rank = len(input_tensor_shape)
@@ -172,48 +188,32 @@ def make_node(
         'optype': graph_node.op,
         'shape': output_tensor_shape,
         'dtype': dtype,
+        'is_dequantized': True,
+        'nhwc': input_nhwc,
     }
 
     # Generation of TF OP
 
-    # Convert w_zero_point and w_scale to 1-D if scalar
-    if len(input_weights_zero_point.shape) == 0:
-        input_weights_zero_point = tf.fill([input_tensor.shape[-1]//group], input_weights_zero_point)
-    elif len(input_weights_zero_point.shape) > 1:
-        error(
-            f'Unsupported zero point: {graph_node.name} {input_weights_zero_point}'
-        )
-        sys.exit(1)
-
-    weights_scale_is_scalar = False
-    if len(input_weights_scale.shape) == 0:
-        weights_scale_is_scalar = True
-        input_weights_scale = tf.fill([input_tensor.shape[-1]//group], input_weights_scale)
-    elif len(input_weights_scale.shape) > 1:
-        error(
-            f'Unsupported scalet: {graph_node.name} {input_weights_scale}'
-        )
-        sys.exit(1)
-
     # Dequantize variables to float32
-    input_tensor = _dequantize_tensor(
-        base=input_tensor,
-        zero_point=input_tensor_zero_point,
-        scale=input_tensor_scale,
-    )
+    if input_is_dequantized:
+        input_tensor = tf.cast(input_tensor, tf.float32)
+    else:
+        input_tensor = _dequantize_tensor(
+            base=input_tensor,
+            zero_point=input_tensor_zero_point,
+            scale=input_tensor_scale,
+        )
     input_weights = _dequantize_weights(
         base=input_weights,
         zero_point=input_weights_zero_point,
         scale=input_weights_scale,
-        scale_is_scalar=weights_scale_is_scalar,
     )
-    y_zero_point = tf.cast(y_zero_point, tf.float32)
 
     # if bias is defined save it here
     if input_bias is not None:
         input_bias = tf.cast(input_bias, tf.float32)
-        input_bias_scale = input_tensor_scale * input_weights_scale
-        input_bias = tf.round(input_bias / input_bias_scale)
+        input_bias_scale = tf.cast(input_tensor_scale, tf.float32) * tf.cast(input_weights_scale, tf.float32)
+        input_bias = input_bias * input_bias_scale
 
     """
     Conv1D
@@ -260,7 +260,7 @@ def make_node(
         depthwise = bool(group == input_tensor_shape[-1])
 
     if depthwise is True:
-        depthwise_filter_shape = list(input_weights_shape[0:2]) + [-1, input_weights_shape[3] // group]
+        depthwise_filter_shape = list(input_weights_shape[0:2]) + [input_weights_shape[2], input_weights_shape[3] // group]
         input_weights = tf.reshape(input_weights, depthwise_filter_shape)
 
     # Conv
@@ -308,27 +308,23 @@ def make_node(
             )
         tf_op_type = tf.nn.depthwise_conv2d
 
-    # Process output
-    scaled_conv_node = tf.add(
-        x=tf.round(
-            tf.divide(
-                x=conv_node,
-                y=y_scale,
-            ),
-        ),
-        y=y_zero_point,
-    )
-
-    # Add bias to the convolution
+    # Add bias to the convolution (float)
     if input_bias is not None:
-        scaled_conv_node = tf.add(
-            x=scaled_conv_node,
+        conv_node = tf.add(
+            x=conv_node,
             y=input_bias,
         )
 
-    casted_conv_node = tf.cast(scaled_conv_node, output_dtype)
+    # quantize then dequantize to float32
+    y_scale = tf.cast(y_scale, tf.float32)
+    y_zero_point = tf.cast(y_zero_point, tf.float32)
+    quantized = tf.round(tf.divide(conv_node, y_scale)) + y_zero_point
+    qmin, qmax = _get_qmin_qmax(output_quant_dtype)
+    if qmin is not None and qmax is not None:
+        quantized = tf.clip_by_value(quantized, qmin, qmax)
+    dequantized = tf.multiply(tf.subtract(quantized, y_zero_point), y_scale)
 
-    tf_layers_dict[graph_node_output.name]['tf_node'] = casted_conv_node
+    tf_layers_dict[graph_node_output.name]['tf_node'] = dequantized
 
     # Generation of Debug Info
     tf_layers_dict[graph_node_output.name]['tf_node_info'] = \
@@ -349,4 +345,3 @@ def make_node(
                 },
             }
         )
-