onnx2tf 1.29.8__py3-none-any.whl → 1.29.10__py3-none-any.whl

onnx2tf/ops/LpPool.py

@@ -0,0 +1,296 @@
+import random
+random.seed(0)
+import numpy as np
+np.random.seed(0)
+import tensorflow as tf
+import onnx_graphsurgeon as gs
+from onnx2tf.utils.common_functions import (
+    get_constant_or_variable,
+    remove_dilations,
+    print_node_info,
+    inverted_operation_enable_disable,
+    make_tf_node_info,
+    get_replacement_parameter,
+    pre_process_transpose,
+    post_process_transpose,
+    calc_tf_pooling_pads,
+    calc_extra_padding_with_ceil,
+    transpose_with_flexing_deterrence,
+)
+from onnx2tf.utils.logging import *
+
+INF_INDEX_VALUE: int = 4294967296
+
+
+def _kernel_size_const(kernel_shape, dtype):
+    if isinstance(kernel_shape, (list, tuple, np.ndarray)):
+        size = 1
+        for k in kernel_shape:
+            size *= int(k)
+        return tf.cast(size, dtype)
+    if tf.is_tensor(kernel_shape):
+        return tf.cast(tf.reduce_prod(kernel_shape), dtype)
+    return tf.cast(int(kernel_shape), dtype)
+
+
+@print_node_info
+@inverted_operation_enable_disable
+@get_replacement_parameter
+def make_node(
+    *,
+    graph_node: gs.Node,
+    tf_layers_dict: dict,
+    **kwargs: dict,
+):
+    """LpPool
+
+    Parameters
+    ----------
+    graph_node: gs.Node
+        graph_surgeon Node
+
+    tf_layers_dict: dict
+        optype, shape, dtype, tensorflow graph
+    """
+    before_op_output_shape_trans_1 = \
+        tf_layers_dict.get(graph_node.inputs[0].name, {}).get('before_op_output_shape_trans', True)
+    before_op_output_shape_trans = \
+        before_op_output_shape_trans_1
+
+    graph_node_input = get_constant_or_variable(
+        graph_node.inputs[0],
+        before_op_output_shape_trans,
+    )
+    graph_node_output: gs.Variable = graph_node.outputs[0]
+    shape = graph_node_output.shape
+    dtype = graph_node_output.dtype
+
+    input_tensor = tf_layers_dict[graph_node_input.name]['tf_node'] \
+        if isinstance(graph_node_input, gs.Variable) else graph_node_input
+    input_tensor_shape = input_tensor.shape
+    input_tensor_rank = len(input_tensor_shape)
+
+    # Pre-process transpose
+    input_tensor = pre_process_transpose(
+        value_before_transpose=input_tensor,
+        param_target='inputs',
+        param_name=graph_node.inputs[0].name,
+        **kwargs,
+    )
+
+    # Workaround to avoid as many conversion failures as possible
+    # for models with useless Transpose immediately before them.
+    # If the input geometry of the ONNX and the input geometry of the TF model match,
+    # the input geometry on the TF model side is forcibly transposed to the NWC or NHWC or NDHWC format.
+    # However, if all dimensions of CW or CHW or CDHW have the same value,
+    # the forced transposition process is skipped because it may destroy the structure of the model.
+    onnx_input_shape = [
+        dim if isinstance(dim, int) else None for dim in graph_node.inputs[0].shape
+    ] if graph_node.inputs[0].shape is not None else None
+    tf_input_shape = [
+        dim if isinstance(dim, int) else None for dim in input_tensor_shape
+    ]
+    if onnx_input_shape is not None \
+        and len(onnx_input_shape) > 1 and len(tf_input_shape) > 1 \
+        and onnx_input_shape == tf_input_shape:
+
+        shape_for_judging_skip = [
+            dim if dim is not None else INF_INDEX_VALUE for dim in onnx_input_shape[1:]
+        ]
+        if shape_for_judging_skip.count(shape_for_judging_skip[0]) != len(shape_for_judging_skip):
+            if len(onnx_input_shape) == 3:
+                # 1D
+                input_tensor = transpose_with_flexing_deterrence(
+                    input_tensor=input_tensor,
+                    perm=[0,2,1],
+                    **kwargs,
+                )
+            elif len(onnx_input_shape) == 4:
+                # 2D
+                input_tensor = transpose_with_flexing_deterrence(
+                    input_tensor=input_tensor,
+                    perm=[0,2,3,1],
+                    **kwargs,
+                )
+            elif len(onnx_input_shape) == 5:
+                # 3D
+                input_tensor = transpose_with_flexing_deterrence(
+                    input_tensor=input_tensor,
+                    perm=[0,2,3,4,1],
+                    **kwargs,
+                )
+
+    auto_pad = graph_node.attrs.get('auto_pad', 'NOTSET')
+    ceil_mode = bool(graph_node.attrs.get('ceil_mode', 0))
+    kernel_shape = graph_node.attrs['kernel_shape']
+    spatial_size = len(kernel_shape)
+    dilations = graph_node.attrs.get('dilations', [1] * spatial_size)
+    pads = graph_node.attrs.get('pads', [0] * spatial_size * 2)
+    strides = graph_node.attrs.get('strides', [1] * spatial_size)
+    p = float(graph_node.attrs.get('p', 2))
+
+    input_tensor_shape = input_tensor.shape.as_list()
+    is_known_shape = None not in input_tensor_shape[1:]
+
+    # default tensorflow action is 'SAME_UPPER' mode (extra padding in the end for odd numbers)
+    # explicit pad layer is added for tensorflow incompatible cases
+    tf_pad_mode = 'VALID'
+    is_explicit_padding = False
+    dilated_kernel_shape = kernel_shape
+    if dilations != [1] * spatial_size:
+        dilated_kernel_shape = [(k - 1) * d for k, d in zip(kernel_shape, dilations)]
+
+    tf_pads = calc_tf_pooling_pads(
+        input_shape=input_tensor_shape,
+        kernel=dilated_kernel_shape,
+        strides=strides,
+        input_tensor=input_tensor,
+    )
+
+    # onnx padding value is ignored if auto_pad is not 'NOTSET'
+    if auto_pad == 'NOTSET':
+
+        # check if onnx padding is same with tensorflow padding mode 'SAME'
+        # this is to avoid flex operations since tflite has no builtin pooling with manual padding value
+        if is_known_shape and pads != [0] * spatial_size * 2 and tf_pads == pads:
+            auto_pad = 'SAME_UPPER'
+            tf_pad_mode = 'SAME'
+
+        else:
+            auto_pad = 'VALID'
+            is_explicit_padding = True
+
+            # extra padding may be needed for ceiling
+            # this padding is added to end side (right, bottom) only
+            if ceil_mode:
+                extra_pads = \
+                    calc_extra_padding_with_ceil(
+                        input_shape=input_tensor_shape[1:-1],
+                        kernel=kernel_shape,
+                        pads=pads,
+                        dilations=dilations,
+                        strides=strides,
+                    )
+                pads = pads[:len(pads) // 2] + [p + e for p, e in zip(pads[len(pads) // 2:], extra_pads)]
+
+            tf_pads = pads
+
+    elif auto_pad == 'SAME_UPPER':
+        tf_pad_mode = 'SAME'
+
+    elif auto_pad == 'SAME_LOWER':
+        is_explicit_padding = True
+
+    elif auto_pad == 'VALID':
+        tf_pads = [0] * spatial_size * 2
+
+    else:
+        error_msg = Color.RED(f'ERROR:') + ' ' + \
+                    f'Wrong auto_pad parameter in LpPool: {auto_pad}.'
+        raise ValueError(error_msg)
+
+    # add extra pad layer if needed
+    if is_explicit_padding and tf_pads != [0] * spatial_size * 2:
+        warn(
+            f'Tensorflow incompatible padding detected. ' \
+            f'Extra pad layer is inserted automatically. '
+        )
+
+        if auto_pad == 'SAME_LOWER':
+            # switch the order of pads
+            tf_pads = [i for tup in zip(tf_pads[len(tf_pads) // 2:], tf_pads[:len(tf_pads) // 2]) for i in tup]
+
+        # convert to tensorflow padding format
+        tf_pads = \
+            [[0, 0]] + \
+            [list(i) for i in zip(tf_pads[:len(tf_pads) // 2], tf_pads[len(tf_pads) // 2:])] + \
+            [[0, 0]]
+
+        padded_tensor = tf.pad(
+            tensor=input_tensor,
+            paddings=tf_pads,
+            mode='CONSTANT',
+            constant_values=0.0,
+        )
+
+    else:
+        padded_tensor = input_tensor
+
+    # Preserving Graph Structure (Dict)
+    tf_layers_dict[graph_node_output.name] = {
+        'optype': graph_node.op,
+        'shape': shape,
+        'dtype': dtype,
+        'nhwc': True,
+    }
+
+    # Generation of TF OP
+    tf_op_type = None
+    abs_p_tensor = tf.pow(tf.abs(padded_tensor), p) if p != 1.0 else tf.abs(padded_tensor)
+    kernel_size_const = _kernel_size_const(kernel_shape, abs_p_tensor.dtype)
+
+    if spatial_size < 4 and (strides == [1] * spatial_size or dilations == [1] * spatial_size):
+        pooled_tensor = tf.nn.pool(
+            input=abs_p_tensor,
+            window_shape=kernel_shape,
+            dilations=dilations,
+            strides=strides,
+            padding=tf_pad_mode.upper(),
+            pooling_type='AVG',
+        )
+        tf_op_type = tf.nn.pool
+    else:
+        # TODO: Dilated LpPool with strides is broken for 3D and above, need to be fixed
+        if spatial_size >= 3:
+            error_msg = Color.RED(f'ERROR:') + ' ' \
+                        f'Dilated LpPool with strides is not supported for 3D and above for now. '
+            print(error_msg)
+            raise NotImplementedError(error_msg)
+
+        abs_p_tensor = remove_dilations(
+            input_tensor=abs_p_tensor,
+            kernel_shape=kernel_shape,
+            spatial_size=spatial_size,
+            strides=strides,
+            dilations=dilations,
+        )
+        pooled_tensor = tf.nn.pool(
+            input=abs_p_tensor,
+            window_shape=kernel_shape,
+            strides=kernel_shape,
+            padding='VALID',
+            pooling_type='AVG',
+        )
+        tf_op_type = tf.nn.pool
+
+    pooled_tensor = pooled_tensor * kernel_size_const
+    tf_layers_dict[graph_node_output.name]['tf_node'] = \
+        tf.pow(pooled_tensor, 1.0 / p) if p != 1.0 else pooled_tensor
+
+    # Post-process transpose
+    tf_layers_dict[graph_node_output.name]['tf_node'] = post_process_transpose(
+        value_before_transpose=tf_layers_dict[graph_node_output.name]['tf_node'],
+        param_target='outputs',
+        param_name=graph_node.outputs[0].name,
+        **kwargs,
+    )
+
+    # Generation of Debug Info
+    tf_layers_dict[graph_node_output.name]['tf_node_info'] = \
+        make_tf_node_info(
+            node_info={
+                'tf_op_type': tf_op_type,
+                'tf_inputs': {
+                    'input': input_tensor,
+                    'kernel_shape': kernel_shape,
+                    'strides': strides,
+                    'dilations': dilations,
+                    'padding': tf_pads if tf_pad_mode != 'same' else tf_pad_mode,
+                    'ceil_mode': ceil_mode,
+                    'p': p,
+                },
+                'tf_outputs': {
+                    'output': tf_layers_dict[graph_node_output.name]['tf_node'],
+                },
+            }
+        )

onnx2tf/ops/MaxRoiPool.py

@@ -0,0 +1,236 @@
+import random
+random.seed(0)
+import numpy as np
+np.random.seed(0)
+import tensorflow as tf
+import onnx_graphsurgeon as gs
+from onnx2tf.utils.common_functions import (
+    get_constant_or_variable,
+    print_node_info,
+    inverted_operation_enable_disable,
+    make_tf_node_info,
+    get_replacement_parameter,
+    pre_process_transpose,
+    post_process_transpose,
+    transpose_with_flexing_deterrence,
+)
+from onnx2tf.utils.logging import *
+
+INF_INDEX_VALUE: int = 4294967296
+
+
+@print_node_info
+@inverted_operation_enable_disable
+@get_replacement_parameter
+def make_node(
+    *,
+    graph_node: gs.Node,
+    tf_layers_dict: dict,
+    **kwargs: dict,
+):
+    """MaxRoiPool
+
+    Parameters
+    ----------
+    graph_node: gs.Node
+        graph_surgeon Node
+
+    tf_layers_dict: dict
+        optype, shape, dtype, tensorflow graph
+    """
+    before_op_output_shape_trans_1 = \
+        tf_layers_dict.get(graph_node.inputs[0].name, {}).get('before_op_output_shape_trans', True)
+    before_op_output_shape_trans_2 = \
+        tf_layers_dict.get(graph_node.inputs[1].name, {}).get('before_op_output_shape_trans', True)
+    before_op_output_shape_trans = \
+        before_op_output_shape_trans_1 and before_op_output_shape_trans_2
+
+    graph_node_input_1 = get_constant_or_variable(
+        graph_node.inputs[0],
+        before_op_output_shape_trans,
+    )
+    graph_node_input_2 = get_constant_or_variable(
+        graph_node.inputs[1],
+        before_op_output_shape_trans,
+    )
+    graph_node_output: gs.Variable = graph_node.outputs[0]
+    shape = graph_node_output.shape
+    dtype = graph_node_output.dtype
+
+    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
+    rois = tf_layers_dict[graph_node_input_2.name]['tf_node'] \
+        if isinstance(graph_node_input_2, gs.Variable) else graph_node_input_2
+
+    # Pre-process transpose
+    input_tensor = pre_process_transpose(
+        value_before_transpose=input_tensor,
+        param_target='inputs',
+        param_name=graph_node.inputs[0].name,
+        **kwargs,
+    )
+
+    # Workaround to avoid as many conversion failures as possible
+    # for models with useless Transpose immediately before them.
+    # If the input geometry of the ONNX and the input geometry of the TF model match,
+    # the input geometry on the TF model side is forcibly transposed to the NHWC format.
+    # However, if all dimensions of CHW have the same value,
+    # the forced transposition process is skipped because it may destroy the structure of the model.
+    onnx_input_shape = [
+        dim if isinstance(dim, int) else None for dim in graph_node.inputs[0].shape
+    ] if graph_node.inputs[0].shape is not None else None
+    tf_input_shape = [
+        dim if isinstance(dim, int) else None for dim in input_tensor.shape
+    ]
+    if onnx_input_shape is not None \
+        and len(onnx_input_shape) > 1 and len(tf_input_shape) > 1 \
+        and onnx_input_shape == tf_input_shape:
+
+        shape_for_judging_skip = [
+            dim if dim is not None else INF_INDEX_VALUE for dim in onnx_input_shape[1:]
+        ]
+        if shape_for_judging_skip.count(shape_for_judging_skip[0]) != len(shape_for_judging_skip):
+            if len(onnx_input_shape) == 4:
+                # 2D
+                input_tensor = transpose_with_flexing_deterrence(
+                    input_tensor=input_tensor,
+                    perm=[0,2,3,1],
+                    **kwargs,
+                )
+
+    pooled_shape = graph_node.attrs.get('pooled_shape', None)
+    if pooled_shape is None or len(pooled_shape) != 2:
+        error_msg = \
+            Color.RED(f'ERROR:') + ' ' + \
+            f'pooled_shape is required for MaxRoiPool. ' \
+            f'graph_node.name: {graph_node.name}, pooled_shape: {pooled_shape}'
+        print(error_msg)
+        raise ValueError(error_msg)
+
+    pooled_h = int(pooled_shape[0])
+    pooled_w = int(pooled_shape[1])
+    spatial_scale = float(graph_node.attrs.get('spatial_scale', 1.0))
+
+    # Preserving Graph Structure (Dict)
+    tf_layers_dict[graph_node_output.name] = {
+        'optype': graph_node.op,
+        'shape': shape,
+        'dtype': dtype,
+        'nhwc': True,
+    }
+
+    # Generation of TF OP
+    rois = tf.cast(rois, tf.float32)
+    if rois.shape.rank == 1:
+        rois = tf.expand_dims(rois, axis=0)
+
+    channels_static = input_tensor.shape[-1]
+    channel_spec = tf.TensorSpec(
+        shape=(channels_static,) if channels_static is not None else (None,),
+        dtype=input_tensor.dtype,
+    )
+    row_spec = tf.TensorSpec(
+        shape=(pooled_w, channels_static) if channels_static is not None else (pooled_w, None),
+        dtype=input_tensor.dtype,
+    )
+    roi_spec = tf.TensorSpec(
+        shape=(pooled_h, pooled_w, channels_static) if channels_static is not None else (pooled_h, pooled_w, None),
+        dtype=input_tensor.dtype,
+    )
+
+    def roi_pool_single(roi):
+        batch_idx = tf.cast(roi[0], tf.int32)
+        x1, y1, x2, y2 = tf.unstack(roi[1:5])
+        x1 = x1 * spatial_scale
+        y1 = y1 * spatial_scale
+        x2 = x2 * spatial_scale
+        y2 = y2 * spatial_scale
+
+        roi_start_w = tf.cast(tf.round(x1), tf.int32)
+        roi_start_h = tf.cast(tf.round(y1), tf.int32)
+        roi_end_w = tf.cast(tf.round(x2), tf.int32)
+        roi_end_h = tf.cast(tf.round(y2), tf.int32)
+
+        height = tf.shape(input_tensor)[1]
+        width = tf.shape(input_tensor)[2]
+
+        roi_start_w = tf.clip_by_value(roi_start_w, 0, width)
+        roi_start_h = tf.clip_by_value(roi_start_h, 0, height)
+        roi_end_w = tf.clip_by_value(roi_end_w, 0, width)
+        roi_end_h = tf.clip_by_value(roi_end_h, 0, height)
+
+        roi_width = tf.maximum(roi_end_w - roi_start_w + 1, 1)
+        roi_height = tf.maximum(roi_end_h - roi_start_h + 1, 1)
+
+        bin_size_h = tf.cast(roi_height, tf.float32) / tf.cast(pooled_h, tf.float32)
+        bin_size_w = tf.cast(roi_width, tf.float32) / tf.cast(pooled_w, tf.float32)
+
+        channels_dynamic = tf.shape(input_tensor)[-1]
+        zero = tf.zeros([channels_dynamic], dtype=input_tensor.dtype)
+
+        def pool_bin(ph, pw):
+            ph_f = tf.cast(ph, tf.float32)
+            pw_f = tf.cast(pw, tf.float32)
+            hstart = tf.cast(tf.floor(ph_f * bin_size_h), tf.int32) + roi_start_h
+            hend = tf.cast(tf.ceil((ph_f + 1.0) * bin_size_h), tf.int32) + roi_start_h
+            wstart = tf.cast(tf.floor(pw_f * bin_size_w), tf.int32) + roi_start_w
+            wend = tf.cast(tf.ceil((pw_f + 1.0) * bin_size_w), tf.int32) + roi_start_w
+
+            hstart = tf.clip_by_value(hstart, 0, height)
+            hend = tf.clip_by_value(hend, 0, height)
+            wstart = tf.clip_by_value(wstart, 0, width)
+            wend = tf.clip_by_value(wend, 0, width)
+
+            is_empty = tf.logical_or(hend <= hstart, wend <= wstart)
+
+            def do_max():
+                region = input_tensor[batch_idx, hstart:hend, wstart:wend, :]
+                return tf.reduce_max(region, axis=[0,1])
+
+            return tf.cond(is_empty, lambda: zero, do_max)
+
+        def pool_row(ph):
+            return tf.map_fn(
+                lambda pw: pool_bin(ph, pw),
+                tf.range(pooled_w),
+                fn_output_signature=channel_spec,
+            )
+
+        return tf.map_fn(
+            pool_row,
+            tf.range(pooled_h),
+            fn_output_signature=row_spec,
+        )
+
+    pooled_tensor = tf.map_fn(
+        roi_pool_single,
+        rois,
+        fn_output_signature=roi_spec,
+    )
+
+    tf_layers_dict[graph_node_output.name]['tf_node'] = pooled_tensor
+
+    # Post-process transpose
+    tf_layers_dict[graph_node_output.name]['tf_node'] = post_process_transpose(
+        value_before_transpose=tf_layers_dict[graph_node_output.name]['tf_node'],
+        param_target='outputs',
+        param_name=graph_node.outputs[0].name,
+        **kwargs,
+    )
+
+    # Generation of Debug Info
+    tf_layers_dict[graph_node_output.name]['tf_node_info'] = \
+        make_tf_node_info(
+            node_info={
+                'tf_op_type': 'MaxRoiPool',
+                'tf_inputs': {
+                    'input': input_tensor,
+                    'rois': rois,
+                    'pooled_shape': pooled_shape,
+                    'spatial_scale': spatial_scale,
+                },
+                'tf_outputs': {
+                    'output': tf_layers_dict[graph_node_output.name]['tf_node'],
+                },
+            }
+        )

onnx2tf/ops/Unsqueeze.py

@@ -69,52 +69,58 @@ def make_node(
         if isinstance(graph_node_input_1, gs.Variable) \
             and 'nhwc' in tf_layers_dict[graph_node_input_1.name].keys() else False
 
+    input_tensor_shape = None
+    tensor_rank = None
     if input_tensor.shape != tf.TensorShape(None):
         input_tensor_shape = list(input_tensor.shape)
-        tensor_rank = len(input_tensor_shape)
-    elif graph_node_output.shape is not None:
+    elif graph_node_output.shape is not None and axes is not None:
         input_tensor_shape = [
             dim for idx, dim in enumerate(graph_node_output.shape) if idx not in axes
         ]
         input_tensor_shape = [
             dim if not isinstance(dim, str) else None for dim in input_tensor_shape
         ]
+    if input_tensor_shape is not None:
         tensor_rank = len(input_tensor_shape)
 
     if isinstance(axes, list) or (isinstance(axes, np.ndarray) and len(axes.shape) > 0):
-        if nhwc:
-            axes = [
-                convert_axis(
-                    axis=idx,
-                    tensor_rank=tensor_rank+len(axes),
-                    before_op_output_shape_trans=True,
-                ) for idx in axes
-            ]
-        elif not nhwc and (isinstance(axes, list) and len(axes) == 1 or isinstance(axes, np.ndarray) and len(axes.shape) == 1) and axes[0] == -1:
-            axes = [
-                convert_axis(
-                    axis=idx,
-                    tensor_rank=tensor_rank+len(axes),
-                    before_op_output_shape_trans=before_op_output_shape_trans,
-                ) for idx in axes
-            ]
+        if tensor_rank is not None:
+            if nhwc:
+                axes = [
+                    convert_axis(
+                        axis=idx,
+                        tensor_rank=tensor_rank+len(axes),
+                        before_op_output_shape_trans=True,
+                    ) for idx in axes
+                ]
+            elif not nhwc and (isinstance(axes, list) and len(axes) == 1 or isinstance(axes, np.ndarray) and len(axes.shape) == 1) and axes[0] == -1:
+                axes = [
+                    convert_axis(
+                        axis=idx,
+                        tensor_rank=tensor_rank+len(axes),
+                        before_op_output_shape_trans=before_op_output_shape_trans,
+                    ) for idx in axes
+                ]
+            else:
+                axes = [idx for idx in axes]
         else:
-            axes = [idx for idx in axes]
+            axes = [int(idx) for idx in axes]
     elif axes is not None and isinstance(axes, np.ndarray) and len(axes.shape) == 0:
-        if nhwc:
-            axes = convert_axis(
-                axis=axes,
-                tensor_rank=tensor_rank+1,
-                before_op_output_shape_trans=True,
-            )
-        elif not nhwc and (isinstance(axes, list) and len(axes) == 1 or isinstance(axes, np.ndarray) and len(axes.shape) == 1) and axes[0] == -1:
-            axes = [
-                convert_axis(
-                    axis=idx,
-                    tensor_rank=tensor_rank+len(axes),
-                    before_op_output_shape_trans=before_op_output_shape_trans,
-                ) for idx in axes
-            ]
+        if tensor_rank is not None:
+            if nhwc:
+                axes = convert_axis(
+                    axis=axes,
+                    tensor_rank=tensor_rank+1,
+                    before_op_output_shape_trans=True,
+                )
+            elif not nhwc and (isinstance(axes, list) and len(axes) == 1 or isinstance(axes, np.ndarray) and len(axes.shape) == 1) and axes[0] == -1:
+                axes = [
+                    convert_axis(
+                        axis=idx,
+                        tensor_rank=tensor_rank+len(axes),
+                        before_op_output_shape_trans=before_op_output_shape_trans,
+                    ) for idx in axes
+                ]
         axes = list(axes[np.newaxis])
 
     if axes is not None and isinstance(axes, list) and len(axes) > 0:
@@ -128,11 +134,13 @@ def make_node(
         **kwargs,
     )
 
-    new_shape = copy.deepcopy(input_tensor_shape)
-    for idx in axes:
-        new_shape.insert(idx, 1)
+    new_shape = None
+    if input_tensor_shape is not None and axes is not None:
+        new_shape = copy.deepcopy(input_tensor_shape)
+        for idx in axes:
+            new_shape.insert(idx, 1)
 
-    new_shape = [dim if dim is not None else -1 for dim in new_shape]
+        new_shape = [dim if dim is not None else -1 for dim in new_shape]
 
     # Preserving Graph Structure (Dict)
     tf_layers_dict[graph_node_output.name] = {
@@ -234,6 +242,29 @@ def make_node(
             tf.identity(input=input_tensor)
         tf_type = tf.identity
 
+    elif not shape_replaced \
+        and new_shape is None:
+        axes_list = axes
+        if axes_list is None:
+            axes_list = []
+        elif isinstance(axes_list, np.ndarray):
+            axes_list = axes_list.tolist() if axes_list.shape != () else [int(axes_list)]
+        elif not isinstance(axes_list, list):
+            axes_list = [int(axes_list)]
+
+        unsqueeze_tensor = input_tensor
+        for axis_idx, axis in enumerate(axes_list):
+            axis_val = int(axis) if isinstance(axis, (np.integer, np.int64, np.int32)) else axis
+            if isinstance(axis_val, int) and axis_val < 0:
+                axis_val = tf.rank(unsqueeze_tensor) + axis_val + 1
+            unsqueeze_tensor = tf.expand_dims(
+                input=unsqueeze_tensor,
+                axis=axis_val,
+                name=graph_node.name if axis_idx == len(axes_list) - 1 else None,
+            )
+        tf_layers_dict[graph_node_output.name]['tf_node'] = unsqueeze_tensor
+        tf_type = tf.expand_dims
+
     elif not shape_replaced \
         and nhwc \
         and len(axes) == 1 \
@@ -247,6 +278,7 @@ def make_node(
         tf_type = tf.expand_dims
 
     elif not shape_replaced \
+        and new_shape is not None \
         and len(new_shape) >= 2 \
         and len([dim for dim in new_shape if dim is None or dim == -1]) >= 2 \
         and not isinstance(axes, int) \

onnx2tf/utils/common_functions.py

@@ -276,6 +276,8 @@ def make_tf_node_info(**kwargs):
 def print_node_info(func):
     @wraps(func)
     def print_wrapper_func(*args, **kwargs):
+        if kwargs.get('suppress_log', False):
+            return func(*args, **kwargs)
         input_onnx_file_path: str = kwargs.get('input_onnx_file_path', None)
         graph_input: gs.Variable = kwargs.get('graph_input', None)
         graph_node: gs.Variable = kwargs.get('graph_node', None)
@@ -4051,6 +4053,7 @@ def dummy_tf_inference(
             for idx, dim in enumerate(input_size):
                 if idx == 0 and input_sizes[0][0] is not None \
                     and len(input_sizes[0]) == len(input_size) \
+                    and len(input_size) >= 2 \
                     and dim is None:
                     # Batch size assignment for input OPs
                     new_input_size.append(input_sizes[0][0])