onnx2tf 1.29.15__py3-none-any.whl → 1.29.17__py3-none-any.whl

onnx2tf/__init__.py

@@ -1,3 +1,3 @@
 from onnx2tf.onnx2tf import convert, main
 
-__version__ = '1.29.15'
+__version__ = '1.29.17'

onnx2tf/onnx2tf.py

@@ -62,6 +62,146 @@ from onnx2tf.utils.enums import (
 from onnx2tf.utils.logging import *
 from sng4onnx import generate as op_name_auto_generate
 
+def fuse_expanded_qdq_to_qdq(
+    *,
+    graph: gs.Graph,
+):
+    def _get_const_value(tensor):
+        if isinstance(tensor, gs.Constant):
+            return tensor.values
+        if isinstance(tensor, gs.Variable) and len(tensor.inputs) == 1:
+            producer = tensor.inputs[0]
+            if producer.op == 'Constant' and 'value' in producer.attrs:
+                return producer.attrs['value'].values
+        return None
+
+    def _split_const_and_var(inputs):
+        if len(inputs) != 2:
+            return None, None
+        const_val = _get_const_value(inputs[0])
+        if const_val is not None:
+            return const_val, inputs[1]
+        const_val = _get_const_value(inputs[1])
+        if const_val is not None:
+            return const_val, inputs[0]
+        return None, None
+
+    nodes_to_remove = []
+    nodes_to_add = []
+
+    for round_node in list(graph.nodes):
+        if round_node.op != 'Round' or len(round_node.inputs) < 1:
+            continue
+
+        round_in = round_node.inputs[0]
+        if len(round_in.inputs) != 1:
+            continue
+        mul1_node = round_in.inputs[0]
+        if mul1_node.op != 'Mul':
+            continue
+        if len(mul1_node.outputs) != 1 or len(mul1_node.outputs[0].outputs) != 1:
+            continue
+
+        inv_scale, x = _split_const_and_var(mul1_node.inputs)
+        if inv_scale is None or x is None:
+            continue
+
+        relu_node = round_node.outputs[0].outputs[0] if round_node.outputs else None
+        if relu_node is None:
+            continue
+        if relu_node.op == 'Relu':
+            relu_out = relu_node.outputs[0]
+        elif relu_node.op in ['Max', 'Maximum']:
+            max_const, max_var = _split_const_and_var(relu_node.inputs)
+            if max_const is None or max_var != round_node.outputs[0]:
+                continue
+            if np.asarray(max_const).size != 1 or float(np.asarray(max_const).item()) != 0.0:
+                continue
+            relu_out = relu_node.outputs[0]
+        else:
+            continue
+
+        if len(relu_out.outputs) != 1:
+            continue
+        min_node = relu_out.outputs[0]
+        if min_node.op not in ['Min', 'Minimum']:
+            continue
+
+        qmax, min_var = _split_const_and_var(min_node.inputs)
+        if qmax is None or min_var != relu_out:
+            continue
+        if np.asarray(qmax).size != 1:
+            continue
+
+        if len(min_node.outputs) != 1 or len(min_node.outputs[0].outputs) != 1:
+            continue
+        mul2_node = min_node.outputs[0].outputs[0]
+        if mul2_node.op != 'Mul':
+            continue
+
+        scale, min_out = _split_const_and_var(mul2_node.inputs)
+        if scale is None or min_out != min_node.outputs[0]:
+            continue
+        if np.asarray(scale).size != 1:
+            continue
+
+        scale_val = float(np.asarray(scale).item())
+        inv_scale_val = float(np.asarray(inv_scale).item())
+        if scale_val == 0.0 or not np.isfinite(scale_val) or not np.isfinite(inv_scale_val):
+            continue
+        if not np.isclose(scale_val * inv_scale_val, 1.0, rtol=1e-3, atol=1e-6):
+            continue
+
+        if len(mul2_node.outputs) != 1:
+            continue
+        output_var = mul2_node.outputs[0]
+
+        # Require linear chain
+        chain_nodes = [mul1_node, round_node, relu_node, min_node, mul2_node]
+        if any(len(n.outputs) == 0 for n in chain_nodes):
+            continue
+        if len(round_node.outputs[0].outputs) != 1 or len(relu_out.outputs) != 1 or len(min_node.outputs[0].outputs) != 1:
+            continue
+
+        # Build QDQ
+        scale_const = gs.Constant(
+            name=f"{mul2_node.name}_scale",
+            values=np.asarray(scale_val, dtype=np.float32),
+        )
+        zero_const = gs.Constant(
+            name=f"{mul2_node.name}_zero_point",
+            values=np.asarray(0, dtype=np.uint8),
+        )
+        quant_out = gs.Variable(
+            name=f"{output_var.name}_quant",
+            dtype=np.uint8,
+            shape=output_var.shape,
+        )
+        q_node = gs.Node(
+            op="QuantizeLinear",
+            name=f"{mul2_node.name}_QuantizeLinear",
+            inputs=[x, scale_const, zero_const],
+            outputs=[quant_out],
+        )
+        dq_node = gs.Node(
+            op="DequantizeLinear",
+            name=f"{mul2_node.name}_DequantizeLinear",
+            inputs=[quant_out, scale_const, zero_const],
+            outputs=[output_var],
+        )
+        output_var.inputs = [dq_node]
+
+        nodes_to_add.extend([q_node, dq_node])
+        nodes_to_remove.extend(chain_nodes)
+
+    if nodes_to_add:
+        graph.nodes.extend(nodes_to_add)
+    if nodes_to_remove:
+        for n in nodes_to_remove:
+            if n in graph.nodes:
+                graph.nodes.remove(n)
+        graph.cleanup().toposort()
+
 def apply_nonzero_passthrough(
     *,
     graph: gs.Graph,
@@ -848,6 +988,7 @@ def convert(
     if hasattr(onnx_graph, 'metadata_props'):
         metadata_props = onnx_graph.metadata_props
     graph = gs.import_onnx(onnx_graph)
+    fuse_expanded_qdq_to_qdq(graph=graph)
 
     # Cut the ONNX graph when an input name is specified that interrupts the conversion
     if not input_names_to_interrupt_model_conversion:

onnx2tf/ops/Add.py

@@ -21,6 +21,7 @@ from onnx2tf.utils.common_functions import (
     disable_unnecessary_transpose,
     shape_unmatched_special_avoidance_workaround,
     merge_two_consecutive_identical_ops_into_one,
+    transpose_with_flexing_deterrence,
     deterring_shape_corruption_due_to_broadcast,
     acquisition_of_validation_data,
     onnx_tf_tensor_validation,
@@ -297,6 +298,117 @@ def make_node(
             )
         tf_type = tf.identity
 
+    def _normalize_dim(dim):
+        return int(dim) if isinstance(dim, (int, np.integer)) else None
+
+    def _get_static_shape(tensor):
+        shape = getattr(tensor, 'shape', None)
+        if shape is None or shape == tf.TensorShape(None):
+            return None
+        return [_normalize_dim(dim) for dim in list(shape)]
+
+    def _shape_match_with_none(expected, actual):
+        if expected is None or actual is None:
+            return False
+        if len(expected) != len(actual):
+            return False
+        for e_dim, a_dim in zip(expected, actual):
+            e_dim = _normalize_dim(e_dim)
+            a_dim = _normalize_dim(a_dim)
+            if e_dim is None or a_dim is None:
+                continue
+            if e_dim != a_dim:
+                return False
+        return True
+
+    def _perm_shape(shape, perm):
+        return [shape[i] for i in perm] if shape is not None else None
+
+    def _limited_perms(rank):
+        identity = list(range(rank))
+        perms = [identity]
+        if rank == 3:
+            perms.append([0, 2, 1])
+        elif rank == 4:
+            perms.extend([[0, 2, 3, 1], [0, 3, 1, 2]])
+        elif rank == 5:
+            perms.extend([[0, 2, 3, 4, 1], [0, 4, 1, 2, 3]])
+        return perms
+
+    def _ranked_perms(perms, input_shape, onnx_shape):
+        if input_shape is None or onnx_shape is None:
+            return perms
+        scored = []
+        for perm in perms:
+            score = 0
+            for out_idx, in_idx in enumerate(perm):
+                if out_idx >= len(onnx_shape) or in_idx >= len(input_shape):
+                    continue
+                o_dim = _normalize_dim(onnx_shape[out_idx])
+                i_dim = input_shape[in_idx]
+                if isinstance(o_dim, int) and isinstance(i_dim, int) and o_dim == i_dim:
+                    score += o_dim
+            scored.append((score, 1 if perm == list(range(len(perm))) else 0, perm))
+        scored.sort(key=lambda x: (x[0], x[1]), reverse=True)
+        return [p for _, _, p in scored]
+
+    # Rescue guard for unexpected broadcasted shapes
+    if not enable_gelu:
+        expected_shape = None
+        if graph_node_output_shape is not None:
+            expected_shape = [_normalize_dim(dim) for dim in list(graph_node_output_shape)]
+        output_shape = _get_static_shape(tf_layers_dict[graph_node_output.name]['tf_node'])
+        input_shape_1 = _get_static_shape(input_tensor_1)
+        input_shape_2 = _get_static_shape(input_tensor_2)
+        if expected_shape is not None \
+            and output_shape is not None \
+            and not _shape_match_with_none(expected_shape, output_shape) \
+            and input_shape_1 is not None \
+            and input_shape_2 is not None \
+            and len(input_shape_1) == len(expected_shape) \
+            and len(input_shape_2) == len(expected_shape):
+
+            rank = len(expected_shape)
+            perms = _limited_perms(rank)
+            perm_list_1 = _ranked_perms(perms, input_shape_1, expected_shape)
+            perm_list_2 = _ranked_perms(perms, input_shape_2, expected_shape)
+            rescue_done = False
+            for perm_1 in perm_list_1:
+                for perm_2 in perm_list_2:
+                    try_input_1 = transpose_with_flexing_deterrence(
+                        input_tensor=input_tensor_1,
+                        perm=perm_1,
+                        **kwargs,
+                    )
+                    try_input_2 = transpose_with_flexing_deterrence(
+                        input_tensor=input_tensor_2,
+                        perm=perm_2,
+                        **kwargs,
+                    )
+                    try:
+                        rescue_tensor = tf.math.add(
+                            x=try_input_1 \
+                                if not isinstance(try_input_1, np.ndarray) \
+                                    else tf.convert_to_tensor(try_input_1),
+                            y=try_input_2 \
+                                if not isinstance(try_input_2, np.ndarray) \
+                                    else tf.convert_to_tensor(try_input_2),
+                            name=graph_node.name,
+                        )
+                    except Exception as ex:
+                        continue
+
+                    rescue_shape = _get_static_shape(rescue_tensor)
+                    if _shape_match_with_none(expected_shape, rescue_shape):
+                        input_tensor_1 = try_input_1
+                        input_tensor_2 = try_input_2
+                        tf_layers_dict[graph_node_output.name]['tf_node'] = rescue_tensor
+                        tf_type = tf.math.add
+                        rescue_done = True
+                        break
+                if rescue_done:
+                    break
+
     # Post-process transpose
     tf_layers_dict[graph_node_output.name]['tf_node'] = \
         post_process_transpose(

onnx2tf/ops/Concat.py

@@ -234,6 +234,31 @@ def make_node(
             and len(value.shape) > 0 else tf.reshape(value, [1]) for value in values
     ]
 
+    def _infer_concat_axis_runtime(values, fallback_axis):
+        if not values:
+            return fallback_axis
+        shapes = [tf.shape(v) for v in values]
+        shapes = tf.stack(shapes)
+        equal_mask = tf.reduce_all(tf.equal(shapes, shapes[0]), axis=0)
+        diff_mask = tf.cast(tf.logical_not(equal_mask), tf.int32)
+        candidate_count = tf.reduce_sum(diff_mask)
+        axis_from_diff = tf.argmax(diff_mask, axis=0, output_type=tf.int32)
+        fallback_axis_tensor = tf.cast(fallback_axis, tf.int32)
+        is_single = tf.cast(tf.equal(candidate_count, 1), tf.int32)
+        return axis_from_diff * is_single + fallback_axis_tensor * (1 - is_single)
+
+    axis_is_dynamic = False
+    if len(values) > 0:
+        all_none = True
+        for value in values:
+            if value.shape is not None and value.shape != tf.TensorShape(None):
+                if not all([s is None for s in value.shape]):
+                    all_none = False
+                    break
+        if all_none:
+            axis_is_dynamic = True
+    axis_for_concat = _infer_concat_axis_runtime(values, axis) if axis_is_dynamic else axis
+
     # Generation of TF OP
     tf_type = None
     if simple_resize:
@@ -266,96 +291,243 @@ def make_node(
         tf_type = tf.constant
 
     else:
+        def _normalize_dim(dim):
+            return int(dim) if isinstance(dim, (int, np.integer)) else None
+
+        def _get_static_shape(tensor):
+            shape = getattr(tensor, 'shape', None)
+            if shape is None or shape == tf.TensorShape(None):
+                return None
+            return [_normalize_dim(dim) for dim in list(shape)]
+
+        def _shape_match_with_none(onnx_shape, tf_shape):
+            if onnx_shape is None or tf_shape is None:
+                return False
+            if len(onnx_shape) != len(tf_shape):
+                return False
+            for o_dim, t_dim in zip(onnx_shape, tf_shape):
+                o_dim = _normalize_dim(o_dim)
+                t_dim = _normalize_dim(t_dim)
+                if o_dim is None or t_dim is None:
+                    continue
+                if o_dim != t_dim:
+                    return False
+            return True
+
+        def _can_concat_shapes(shapes, axis):
+            if shapes is None or any(s is None for s in shapes):
+                return True
+            rank = len(shapes[0])
+            for idx in range(rank):
+                if idx == axis:
+                    continue
+                dims = [s[idx] for s in shapes]
+                known = [d for d in dims if isinstance(d, int)]
+                if len(known) >= 2 and len(set(known)) != 1:
+                    return False
+            return True
+
+        def _perm_shape(shape, perm):
+            return [shape[i] for i in perm] if shape is not None else None
+
+        def _limited_perms(rank):
+            identity = list(range(rank))
+            perms = [identity]
+            if rank == 3:
+                perms.append([0, 2, 1])
+            elif rank == 4:
+                perms.extend([[0, 2, 3, 1], [0, 3, 1, 2]])
+            elif rank == 5:
+                perms.extend([[0, 2, 3, 4, 1], [0, 4, 1, 2, 3]])
+            return perms
+
+        def _base_perms(rank):
+            if rank <= 1:
+                return [list(range(rank))]
+            return [list(p) for p in itertools.permutations(range(rank))]
+
+        def _ranked_perms(perms, input_shape, axis, onnx_shape):
+            identity = list(range(len(perms[0]))) if perms else []
+            scored = []
+            for perm in perms:
+                score = 0
+                if input_shape is not None and onnx_shape is not None:
+                    for out_idx, in_idx in enumerate(perm):
+                        if out_idx == axis:
+                            continue
+                        o_dim = _normalize_dim(onnx_shape[out_idx]) if out_idx < len(onnx_shape) else None
+                        i_dim = input_shape[in_idx] if in_idx < len(input_shape) else None
+                        if isinstance(o_dim, int) and isinstance(i_dim, int) and o_dim == i_dim:
+                            score += o_dim
+                scored.append((score, 1 if perm == identity else 0, perm))
+            scored.sort(key=lambda x: (x[0], x[1]), reverse=True)
+            return [p for _, _, p in scored]
+
         try:
             # normal concat attempt
             tf_layers_dict[graph_node_output.name]['tf_node'] = \
                 tf.concat(
                     values=values,
-                    axis=axis,
+                    axis=axis_for_concat,
                     name=graph_node.name,
                 )
         except:
             # Workaround to reduce error rate when merging tensors with undefined dimensions
-            try:
-                # Attempts to bind with the axis specified in ONNX
-                tf_layers_dict[graph_node_output.name]['tf_node'] = \
-                    tf.concat(
-                        values=values,
-                        axis=int(graph_node.attrs.get('axis', 0)),
-                        name=graph_node.name,
-                    )
-                axis = int(graph_node.attrs.get('axis', 0))
-            except:
-                # If not successful with the same axis as ONNX, try to combine by other axes
-                # Trial in reverse order from the axis at the end
-                value_rank = len(values[0].shape)
-                succeed = False
-                for idx in reversed(range(value_rank)):
+            original_values = values
+            original_shapes = [_get_static_shape(v) for v in original_values]
+            value_rank = getattr(original_values[0].shape, 'rank', None)
+            if value_rank is None:
+                value_rank = len(original_values[0].shape)
+
+            onnx_shape_list = None
+            if onnx_output_shape is not None:
+                onnx_shape_list = [_normalize_dim(dim) for dim in list(onnx_output_shape)]
+
+            onnx_axis = int(graph_node.attrs.get('axis', 0))
+            onnx_axis = onnx_axis + value_rank if onnx_axis < 0 else onnx_axis
+
+            def _axis_score(axis_idx):
+                if onnx_shape_list is not None and axis_idx < len(onnx_shape_list):
+                    onnx_dim = onnx_shape_list[axis_idx]
+                    if isinstance(onnx_dim, int):
+                        return onnx_dim
+                score = 0
+                for shape in original_shapes:
+                    if shape is None or axis_idx >= len(shape):
+                        continue
+                    dim = shape[axis_idx]
+                    if isinstance(dim, int):
+                        score += dim
+                return score
+
+            axis_candidates = list(range(value_rank))
+            axis_candidates.sort(
+                key=lambda a: (_axis_score(a), 1 if a == onnx_axis else 0),
+                reverse=True,
+            )
+
+            base_perms = _base_perms(value_rank)
+            max_combo = 20000
+            if len(base_perms) ** len(original_values) > max_combo:
+                base_perms = _limited_perms(value_rank)
+
+            succeed = False
+            matched = False
+            chosen_axis = None
+            chosen_values = None
+            chosen_tensor = None
+
+            for axis_idx in axis_candidates:
+                perm_lists = [
+                    _ranked_perms(base_perms, shape, axis_idx, onnx_shape_list)
+                    for shape in original_shapes
+                ]
+                for perm_combo in itertools.product(*perm_lists):
+                    permuted_shapes = [
+                        _perm_shape(shape, perm) for shape, perm in zip(original_shapes, perm_combo)
+                    ]
+                    if not _can_concat_shapes(permuted_shapes, axis_idx):
+                        continue
+                    try_values = [
+                        value if perm == list(range(value_rank)) else
+                        transpose_with_flexing_deterrence(
+                            input_tensor=value,
+                            perm=perm,
+                            **kwargs,
+                        )
+                        for value, perm in zip(original_values, perm_combo)
+                    ]
                     try:
-                        tf_layers_dict[graph_node_output.name]['tf_node'] = \
+                        concat_tensor = \
                             tf.concat(
-                                values=values,
-                                axis=idx,
+                                values=try_values,
+                                axis=axis_idx,
                                 name=graph_node.name,
                             )
-                        axis = idx
+                    except:
+                        continue
+
+                    if not succeed:
                         succeed = True
+                        chosen_axis = axis_idx
+                        chosen_values = try_values
+                        chosen_tensor = concat_tensor
+
+                    if onnx_shape_list is None:
+                        matched = True
+                        chosen_axis = axis_idx
+                        chosen_values = try_values
+                        chosen_tensor = concat_tensor
                         break
-                    except:
-                        pass
-                if not succeed:
-                    raise
+
+                    output_shape = _get_static_shape(concat_tensor)
+                    if _shape_match_with_none(onnx_shape_list, output_shape):
+                        matched = True
+                        chosen_axis = axis_idx
+                        chosen_values = try_values
+                        chosen_tensor = concat_tensor
+                        break
+                if matched:
+                    break
+
+            if succeed:
+                tf_layers_dict[graph_node_output.name]['tf_node'] = chosen_tensor
+                axis = chosen_axis
+                values = chosen_values
+            else:
+                raise
 
         # Attempts to force axis correction when the number of axes in the combined tensor do not exactly match.
         # However, if more than 2 patterns of correct answers exist, give up the correction.
         # This workaround is useful when automatic axis correction is practically difficult,
         # such as when all tensors to be combined originate from Transpose or Reshape.
         # https://github.com/PINTO0309/onnx2tf/issues/473
-        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:
-
-                # Search for valid Concat patterns
-                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=values,
-                                    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
-                    # Review Concat axes only if there is one valid join pattern
-                    if len(matched_axes) == 1:
-                        tf_layers_dict[graph_node_output.name]['tf_node'] = \
-                            tf.concat(
-                                values=values,
-                                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:
+        if not axis_is_dynamic:
+            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:
+
+                    # Search for valid Concat patterns
+                    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=values,
+                                        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
+                        # Review Concat axes only if there is one valid join pattern
+                        if len(matched_axes) == 1:
                             tf_layers_dict[graph_node_output.name]['tf_node'] = \
                                 tf.concat(
                                     values=values,
-                                    axis=onnx_axis,
+                                    axis=matched_axes[0],
                                     name=graph_node.name,
                                 )
-                            axis = onnx_axis
+                            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=values,
+                                        axis=onnx_axis,
+                                        name=graph_node.name,
+                                    )
+                                axis = onnx_axis
 
         # Workaround for post-concat accuracy degradation issues
         # Process only in the case of a Concat of two tensors because the process is too redundant.