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

onnx2tf/__init__.py

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

onnx2tf/ops/ImageDecoder.py

@@ -0,0 +1,147 @@
+import random
+random.seed(0)
+import numpy as np
+np.random.seed(0)
+import tensorflow as tf
+import onnx_graphsurgeon as gs
+import cv2
+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,
+)
+from onnx2tf.utils.logging import *
+
+
+def _as_tensor(value):
+    if isinstance(value, np.ndarray):
+        return tf.convert_to_tensor(value)
+    if isinstance(value, (np.generic, int, float, bool, str, bytes)):
+        return tf.convert_to_tensor(value)
+    return value
+
+
+def _decode_image_np(encoded_stream, pixel_format):
+    if encoded_stream is None:
+        return np.zeros((0, 0, 0), dtype=np.uint8)
+    if encoded_stream.dtype != np.uint8:
+        encoded_stream = encoded_stream.astype(np.uint8)
+    if encoded_stream.size == 0:
+        return np.zeros((0, 0, 0), dtype=np.uint8)
+    if encoded_stream.ndim != 1:
+        encoded_stream = encoded_stream.reshape(-1)
+    try:
+        if pixel_format == 'Grayscale':
+            flag = cv2.IMREAD_GRAYSCALE
+        else:
+            flag = cv2.IMREAD_COLOR
+        decoded = cv2.imdecode(encoded_stream, flag)
+        if decoded is None:
+            raise ValueError('cv2.imdecode failed')
+        if pixel_format == 'RGB':
+            decoded = cv2.cvtColor(decoded, cv2.COLOR_BGR2RGB)
+        if pixel_format == 'Grayscale' and decoded.ndim == 2:
+            decoded = decoded[..., np.newaxis]
+        return decoded.astype(np.uint8)
+    except Exception:
+        if pixel_format == 'Grayscale':
+            return np.zeros((0, 0, 1), dtype=np.uint8)
+        return np.zeros((0, 0, 3), dtype=np.uint8)
+
+
+@print_node_info
+@inverted_operation_enable_disable
+@get_replacement_parameter
+def make_node(
+    *,
+    graph_node: gs.Node,
+    tf_layers_dict: dict,
+    **kwargs: dict,
+):
+    """ImageDecoder
+
+    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
+
+    # Preserving Graph Structure (Dict)
+    tf_layers_dict[graph_node_output.name] = {
+        'optype': graph_node.op,
+        'shape': shape,
+        'dtype': dtype,
+    }
+
+    # Pre-process transpose
+    input_tensor = pre_process_transpose(
+        value_before_transpose=input_tensor,
+        param_target='inputs',
+        param_name=graph_node.inputs[0].name,
+        **kwargs,
+    )
+
+    # Generation of TF OP
+    input_tensor = _as_tensor(input_tensor)
+    pixel_format = graph_node.attrs.get('pixel_format', 'RGB')
+    if pixel_format not in ['RGB', 'BGR', 'Grayscale']:
+        error(
+            f'ImageDecoder pixel_format={pixel_format} is not supported.\n' +
+            f'graph_node.name: {graph_node.name}'
+        )
+        pixel_format = 'RGB'
+
+    decoded = tf.numpy_function(
+        func=lambda x: _decode_image_np(x, pixel_format),
+        inp=[input_tensor],
+        Tout=tf.uint8,
+        name=graph_node.name,
+    )
+    channels = 1 if pixel_format == 'Grayscale' else 3
+    decoded = tf.ensure_shape(decoded, [None, None, channels])
+    tf_layers_dict[graph_node_output.name]['tf_node'] = decoded
+
+    # 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': 'ImageDecoder',
+                'tf_inputs': {
+                    'encoded_stream': input_tensor,
+                    'pixel_format': pixel_format,
+                },
+                'tf_outputs': {
+                    'output': tf_layers_dict[graph_node_output.name]['tf_node'],
+                },
+            }
+        )

onnx2tf/ops/NegativeLogLikelihoodLoss.py

@@ -0,0 +1,237 @@
+import sys
+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,
+    convert_axis,
+    print_node_info,
+    inverted_operation_enable_disable,
+    make_tf_node_info,
+    get_replacement_parameter,
+    pre_process_transpose,
+    post_process_transpose,
+)
+from onnx2tf.utils.enums import NUMPY_DTYPES_TO_TF_DTYPES
+from onnx2tf.utils.logging import *
+
+
+def _as_tensor(value):
+    if isinstance(value, np.ndarray):
+        return tf.convert_to_tensor(value)
+    if isinstance(value, (np.generic, int, float, bool, str, bytes)):
+        return tf.convert_to_tensor(value)
+    return value
+
+
+def _move_class_to_last(tensor, class_axis):
+    rank = tensor.shape.rank
+    if rank is None:
+        rank = tf.rank(tensor)
+    if isinstance(rank, int):
+        if class_axis == rank - 1:
+            return tensor, None
+        perm = [i for i in range(rank) if i != class_axis] + [class_axis]
+        return tf.transpose(tensor, perm=perm), perm
+    return tensor, None
+
+
+@print_node_info
+@inverted_operation_enable_disable
+@get_replacement_parameter
+def make_node(
+    *,
+    graph_node: gs.Node,
+    tf_layers_dict: dict,
+    **kwargs: dict,
+):
+    """NegativeLogLikelihoodLoss
+
+    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
+    if len(graph_node.inputs) >= 3:
+        before_op_output_shape_trans_3 = \
+            tf_layers_dict.get(graph_node.inputs[2].name, {}).get('before_op_output_shape_trans', True)
+        before_op_output_shape_trans = \
+            before_op_output_shape_trans \
+            and before_op_output_shape_trans_3
+
+    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_input_3 = None
+    if len(graph_node.inputs) >= 3:
+        graph_node_input_3 = get_constant_or_variable(
+            graph_node.inputs[2],
+            before_op_output_shape_trans,
+        )
+
+    graph_node_output: gs.Variable = graph_node.outputs[0]
+    shape = graph_node_output.shape
+    dtype = graph_node_output.dtype
+    output_tf_dtype = NUMPY_DTYPES_TO_TF_DTYPES[dtype] \
+        if isinstance(dtype, np.dtype) else 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
+    target_tensor = tf_layers_dict[graph_node_input_2.name]['tf_node'] \
+        if isinstance(graph_node_input_2, gs.Variable) else graph_node_input_2
+    weight_tensor = None
+    if graph_node_input_3 is not None:
+        weight_tensor = tf_layers_dict[graph_node_input_3.name]['tf_node'] \
+            if isinstance(graph_node_input_3, gs.Variable) else graph_node_input_3
+
+    reduction = graph_node.attrs.get('reduction', 'mean')
+    ignore_index = graph_node.attrs.get('ignore_index', None)
+
+    input_rank = len(input_tensor.shape)
+    class_axis = convert_axis(
+        axis=1,
+        tensor_rank=input_rank,
+        before_op_output_shape_trans=before_op_output_shape_trans,
+    )
+
+    # Preserving Graph Structure (Dict)
+    output_entry = {
+        'optype': graph_node.op,
+        'shape': shape,
+        'dtype': dtype,
+    }
+    if reduction == 'none':
+        output_entry['nhwc'] = tf_layers_dict[graph_node_input_1.name]['nhwc'] \
+            if isinstance(graph_node_input_1, gs.Variable) \
+                and 'nhwc' in tf_layers_dict[graph_node_input_1.name].keys() else False
+    tf_layers_dict[graph_node_output.name] = output_entry
+
+    # Pre-process transpose
+    input_tensor = pre_process_transpose(
+        value_before_transpose=input_tensor,
+        param_target='inputs',
+        param_name=graph_node.inputs[0].name,
+        **kwargs,
+    )
+    target_tensor = pre_process_transpose(
+        value_before_transpose=target_tensor,
+        param_target='inputs',
+        param_name=graph_node.inputs[1].name,
+        **kwargs,
+    )
+    if weight_tensor is not None:
+        weight_tensor = pre_process_transpose(
+            value_before_transpose=weight_tensor,
+            param_target='inputs',
+            param_name=graph_node.inputs[2].name,
+            **kwargs,
+        )
+
+    # Generation of TF OP
+    input_tensor = _as_tensor(input_tensor)
+    target_tensor = _as_tensor(target_tensor)
+    if weight_tensor is not None:
+        weight_tensor = _as_tensor(weight_tensor)
+
+    log_prob = input_tensor
+    log_prob, _ = _move_class_to_last(log_prob, class_axis)
+
+    depth = log_prob.shape[-1]
+    if depth is None:
+        depth = tf.shape(log_prob)[-1]
+    depth = tf.cast(depth, tf.int32)
+
+    labels = tf.cast(target_tensor, tf.int32)
+    if ignore_index is not None:
+        ignore_index_val = tf.cast(ignore_index, labels.dtype)
+        mask = tf.equal(labels, ignore_index_val)
+        labels_safe = tf.where(mask, tf.zeros_like(labels), labels)
+    else:
+        mask = None
+        labels_safe = labels
+
+    one_hot = tf.one_hot(
+        indices=labels_safe,
+        depth=depth,
+        axis=-1,
+        dtype=log_prob.dtype,
+    )
+    selected = tf.reduce_sum(log_prob * one_hot, axis=-1)
+    loss = -selected
+
+    weight_per_label = None
+    if weight_tensor is not None:
+        weight_per_label = tf.gather(weight_tensor, labels_safe)
+        weight_per_label = tf.cast(weight_per_label, loss.dtype)
+        if mask is not None:
+            weight_per_label = tf.where(mask, tf.zeros_like(weight_per_label), weight_per_label)
+        loss = loss * weight_per_label
+
+    if mask is not None:
+        loss = tf.where(mask, tf.zeros_like(loss), loss)
+
+    if reduction == 'none':
+        output_tensor = loss
+    elif reduction == 'sum':
+        output_tensor = tf.reduce_sum(loss)
+    elif reduction == 'mean':
+        if weight_per_label is None:
+            output_tensor = tf.reduce_mean(loss)
+        else:
+            denom = tf.reduce_sum(weight_per_label)
+            output_tensor = tf.math.divide_no_nan(tf.reduce_sum(loss), denom)
+    else:
+        error(
+            f'NegativeLogLikelihoodLoss reduction={reduction} is not supported.\n' +
+            f'graph_node.name: {graph_node.name}'
+        )
+        sys.exit(1)
+
+    if output_tf_dtype is not None and output_tensor.dtype != output_tf_dtype:
+        output_tensor = tf.cast(output_tensor, output_tf_dtype)
+
+    tf_layers_dict[graph_node_output.name]['tf_node'] = output_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': 'NegativeLogLikelihoodLoss',
+                'tf_inputs': {
+                    'input': input_tensor,
+                    'target': target_tensor,
+                    'weight': weight_tensor,
+                    'reduction': reduction,
+                    'ignore_index': ignore_index,
+                },
+                'tf_outputs': {
+                    'output': tf_layers_dict[graph_node_output.name]['tf_node'],
+                },
+            }
+        )

onnx2tf/ops/RMSNormalization.py

@@ -0,0 +1,175 @@
+import sys
+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,
+    convert_axis,
+    print_node_info,
+    inverted_operation_enable_disable,
+    make_tf_node_info,
+    get_replacement_parameter,
+    pre_process_transpose,
+    post_process_transpose,
+)
+from onnx2tf.utils.enums import NUMPY_DTYPES_TO_TF_DTYPES
+from onnx2tf.utils.logging import *
+
+
+def _as_tensor(value):
+    if isinstance(value, np.ndarray):
+        return tf.convert_to_tensor(value)
+    if isinstance(value, (np.generic, int, float, bool, str, bytes)):
+        return tf.convert_to_tensor(value)
+    return value
+
+
+@print_node_info
+@inverted_operation_enable_disable
+@get_replacement_parameter
+def make_node(
+    *,
+    graph_node: gs.Node,
+    tf_layers_dict: dict,
+    **kwargs: dict,
+):
+    """RMSNormalization
+
+    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
+    output_tf_dtype = NUMPY_DTYPES_TO_TF_DTYPES[dtype] \
+        if isinstance(dtype, np.dtype) else 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
+    scale_tensor = tf_layers_dict[graph_node_input_2.name]['tf_node'] \
+        if isinstance(graph_node_input_2, gs.Variable) else graph_node_input_2
+
+    # Preserving Graph Structure (Dict)
+    tf_layers_dict[graph_node_output.name] = {
+        'optype': graph_node.op,
+        'shape': shape,
+        'dtype': dtype,
+        'nhwc': tf_layers_dict[graph_node_input_1.name]['nhwc'] \
+            if isinstance(graph_node_input_1, gs.Variable) \
+                and 'nhwc' in tf_layers_dict[graph_node_input_1.name].keys() else False
+    }
+
+    # Pre-process transpose
+    input_tensor = pre_process_transpose(
+        value_before_transpose=input_tensor,
+        param_target='inputs',
+        param_name=graph_node.inputs[0].name,
+        **kwargs,
+    )
+    scale_tensor = pre_process_transpose(
+        value_before_transpose=scale_tensor,
+        param_target='inputs',
+        param_name=graph_node.inputs[1].name,
+        **kwargs,
+    )
+
+    # Generation of TF OP
+    input_tensor = _as_tensor(input_tensor)
+    scale_tensor = _as_tensor(scale_tensor)
+
+    input_rank = input_tensor.shape.rank
+    if input_rank is None and graph_node.inputs[0].shape is not None:
+        input_rank = len(graph_node.inputs[0].shape)
+    if input_rank is None:
+        error(
+            f'RMSNormalization requires known input rank.\n' +
+            f'graph_node.name: {graph_node.name}'
+        )
+        sys.exit(1)
+
+    axis = graph_node.attrs.get('axis', -1)
+    axis = convert_axis(
+        axis=axis,
+        tensor_rank=input_rank,
+        before_op_output_shape_trans=before_op_output_shape_trans,
+    )
+    epsilon = graph_node.attrs.get('epsilon', 1e-05)
+    stash_type = int(graph_node.attrs.get('stash_type', 1))
+
+    axes = list(range(axis, input_rank))
+
+    compute_dtype = input_tensor.dtype
+    if stash_type == 1 and input_tensor.dtype != tf.float32:
+        compute_dtype = tf.float32
+
+    x = tf.cast(input_tensor, compute_dtype)
+    xsquared = tf.math.square(x)
+    xsquared_mean = tf.reduce_mean(xsquared, axis=axes, keepdims=True)
+    rms = tf.sqrt(xsquared_mean + tf.cast(epsilon, compute_dtype))
+    normalized = x / rms
+
+    if compute_dtype != input_tensor.dtype:
+        normalized = tf.cast(normalized, input_tensor.dtype)
+    if scale_tensor.dtype != normalized.dtype:
+        scale_tensor = tf.cast(scale_tensor, normalized.dtype)
+
+    output_tensor = tf.math.multiply(
+        normalized,
+        scale_tensor,
+        name=graph_node.name,
+    )
+
+    if output_tf_dtype is not None and output_tensor.dtype != output_tf_dtype:
+        output_tensor = tf.cast(output_tensor, output_tf_dtype)
+
+    tf_layers_dict[graph_node_output.name]['tf_node'] = output_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': 'RMSNormalization',
+                'tf_inputs': {
+                    'input': input_tensor,
+                    'scale': scale_tensor,
+                    'axis': axis,
+                    'epsilon': epsilon,
+                    'stash_type': stash_type,
+                },
+                'tf_outputs': {
+                    'output': tf_layers_dict[graph_node_output.name]['tf_node'],
+                },
+            }
+        )

onnx2tf/ops/RegexFullMatch.py

@@ -0,0 +1,108 @@
+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,
+)
+
+
+def _as_tensor(value):
+    if isinstance(value, np.ndarray):
+        return tf.convert_to_tensor(value)
+    if isinstance(value, (np.generic, int, float, bool, str, bytes)):
+        return tf.convert_to_tensor(value)
+    return value
+
+
+@print_node_info
+@inverted_operation_enable_disable
+@get_replacement_parameter
+def make_node(
+    *,
+    graph_node: gs.Node,
+    tf_layers_dict: dict,
+    **kwargs: dict,
+):
+    """RegexFullMatch
+
+    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
+
+    # Preserving Graph Structure (Dict)
+    tf_layers_dict[graph_node_output.name] = {
+        'optype': graph_node.op,
+        'shape': shape,
+        'dtype': dtype,
+    }
+
+    # Pre-process transpose
+    input_tensor = pre_process_transpose(
+        value_before_transpose=input_tensor,
+        param_target='inputs',
+        param_name=graph_node.inputs[0].name,
+        **kwargs,
+    )
+
+    # Generation of TF OP
+    input_tensor = _as_tensor(input_tensor)
+    pattern = graph_node.attrs.get('pattern', '')
+
+    tf_layers_dict[graph_node_output.name]['tf_node'] = \
+        tf.strings.regex_full_match(
+            input=input_tensor,
+            pattern=pattern,
+            name=graph_node.name,
+        )
+
+    # 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.strings.regex_full_match,
+                'tf_inputs': {
+                    'input': input_tensor,
+                    'pattern': pattern,
+                },
+                'tf_outputs': {
+                    'output': tf_layers_dict[graph_node_output.name]['tf_node'],
+                },
+            }
+        )