onnx2tf 1.24.0__py3-none-any.whl → 1.25.9__py3-none-any.whl

onnx2tf/__init__.py

@@ -1,3 +1,3 @@
 from onnx2tf.onnx2tf import convert, main
 
-__version__ = '1.24.0'
+__version__ = '1.25.9'

onnx2tf/onnx2tf.py

@@ -25,6 +25,7 @@ os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3'
 os.environ["TF_USE_LEGACY_KERAS"] = '1'
 import tensorflow as tf
 tf.random.set_seed(0)
+from tensorflow.python.saved_model.load import _WrapperFunction
 import tf_keras
 tf_keras.utils.set_random_seed(0)
 tf.config.experimental.enable_op_determinism()
@@ -81,6 +82,7 @@ def convert(
     keep_ncw_or_nchw_or_ncdhw_input_names: Optional[List[str]] = None,
     keep_nwc_or_nhwc_or_ndhwc_input_names: Optional[List[str]] = None,
     keep_shape_absolutely_input_names: Optional[List[str]] = None,
+    input_names_to_interrupt_model_conversion: Optional[List[str]] = None,
     output_names_to_interrupt_model_conversion: Optional[List[str]] = None,
     disable_group_convolution: Optional[bool] = False,
     enable_accumulation_type_float16: Optional[bool] = False,
@@ -285,6 +287,13 @@ def convert(
         e.g.\n
         keep_shape_absolutely_input_names=['input0','input1','input2']
 
+    input_names_to_interrupt_model_conversion: Optional[List[str]]
+        Input names that interrupt model conversion.\n
+        Interrupts model transformation at the specified input name\n
+        and inputs the model partitioned into subgraphs.\n\n
+        e.g.\n
+        input_names_to_interrupt_model_conversion=['input0','input1','input2']
+
     output_names_to_interrupt_model_conversion: Optional[List[str]]
         Output names that interrupt model conversion.\n
         Interrupts model transformation at the specified output name\n
@@ -686,7 +695,77 @@ def convert(
         metadata_props = onnx_graph.metadata_props
     graph = gs.import_onnx(onnx_graph)
 
-    # List Output
+    # Cut the ONNX graph when an input name is specified that interrupts the conversion
+    if not input_names_to_interrupt_model_conversion:
+        input_names = [
+            graph_input.name for graph_input in graph.inputs
+        ]
+    else:
+        try:
+            from sne4onnx import extraction
+        except Exception as ex:
+            error(
+                f'If --input_names_to_interrupt_model_conversion is specified, ' +\
+                f'you must install sne4onnx. pip install sne4onnx'
+            )
+            sys.exit(1)
+        # Cut ONNX graph at specified input position
+        input_names = [
+            input_op_name \
+                for input_op_name in input_names_to_interrupt_model_conversion
+        ]
+        onnx_graph: onnx.ModelProto = \
+            extraction(
+                input_op_names=input_names,
+                output_op_names=[graph_output.name for graph_output in graph.outputs],
+                onnx_graph=onnx_graph,
+            )
+        # Re-import of onnx_graph
+        del graph
+        graph = gs.import_onnx(onnx_graph)
+
+        total_num_nodes = len(graph.nodes)
+        check_count = 0
+        idx = 0
+        while True:
+            # Delete unused nodes
+            if check_count >= total_num_nodes:
+                break
+            op_input_names: List[str] = [inp.name for inp in graph.nodes[idx].inputs]
+            remove_enable = not any(name in input_names for name in op_input_names)
+            if remove_enable:
+                try:
+                    num_input = len(graph.nodes[idx].inputs)
+                    enable_var_input = False
+                    for sub_idx, graph_node_input in enumerate(graph.nodes[idx].inputs):
+                        if isinstance(graph_node_input, gs.Variable):
+                            enable_var_input = True
+                            break
+                        else:
+                            pass
+                    if enable_var_input:
+                        name = graph.nodes[idx].i(sub_idx).name
+                        if any([graph_node.name == name for graph_node in graph.nodes]):
+                            idx += 1
+                        else:
+                            try:
+                                del graph.nodes[idx]
+                            except IndexError:
+                                break
+                    else:
+                        idx += 1
+                except:
+                    try:
+                        del graph.nodes[idx]
+                    except IndexError:
+                        break
+            else:
+                idx += 1
+            check_count += 1
+        onnx_graph = gs.export_onnx(graph=graph, do_type_check=False, **meta_data)
+        if metadata_props is not None:
+            onnx_graph.metadata_props.extend(metadata_props)
+
     # Cut the ONNX graph when an output name is specified that interrupts the conversion
     if not output_names_to_interrupt_model_conversion:
         output_names = [
@@ -706,11 +785,12 @@ def convert(
             output_op_name \
                 for output_op_name in output_names_to_interrupt_model_conversion
         ]
-        onnx_graph: onnx.ModelProto = extraction(
-            input_op_names=[graph_input.name for graph_input in graph.inputs],
-            output_op_names=output_names,
-            onnx_graph=onnx_graph,
-        )
+        onnx_graph: onnx.ModelProto = \
+            extraction(
+                input_op_names=[graph_input.name for graph_input in graph.inputs],
+                output_op_names=output_names,
+                onnx_graph=onnx_graph,
+            )
         # Re-import of onnx_graph
         del graph
         graph = gs.import_onnx(onnx_graph)
@@ -772,6 +852,7 @@ def convert(
             output_name = re.sub('^/', 'wa/', output_name)
             new_output_names.append(output_name)
         output_names = new_output_names
+
     try:
         onnx_graph = gs.export_onnx(graph=graph, do_type_check=False, **meta_data)
         if metadata_props is not None:
@@ -1234,19 +1315,56 @@ def convert(
             # Switch to .pb
             info(Color.GREEN(f'Switch to the output of an optimized protocol buffer file (.pb).'))
         except (KeyError, AssertionError) as e:
+            msg_list = [s for s in e.args if isinstance(s, str)]
+            if len(msg_list) > 0:
+                try:
+                    for s in msg_list:
+                        if 'Failed to add concrete function' in s \
+                            or "Tried to export a function which references an 'untracked' resource" in s:
+                            export_archive = tf_keras.export.ExportArchive()
+                            export_archive.add_endpoint(
+                                name=SIGNATURE_KEY,
+                                fn=lambda *inputs : model(inputs),
+                                input_signature=[tf.TensorSpec(tensor.shape, tensor.dtype, tensor.name) for tensor in model.inputs],
+                            )
+                            export_archive.write_out(output_folder_path)
+                            break
+                except ValueError as e:
+                    msg_list = [s for s in e.args if isinstance(s, str)]
+                    if len(msg_list) > 0:
+                        for s in msg_list:
+                            if 'A root scope name has to match the following pattern' in s:
+                                error(
+                                    f'Generation of saved_model failed because the OP name does not match the following pattern. ^[A-Za-z0-9.][A-Za-z0-9_.\\\\/>-]*$'
+                                )
+                                matches = re.findall(r"'([^']*)'", s)
+                                error(f'{matches[0]}')
+                                error(
+                                    f'Please convert again with the `-osd` or `--output_signaturedefs` option.'
+                                )
+                                sys.exit(1)
+                    else:
+                        error(e)
+                        import traceback
+                        error(traceback.format_exc(), prefix=False)
+            else:
+                error(e)
+                import traceback
+                error(traceback.format_exc(), prefix=False)
+        except ValueError as e:
             msg_list = [s for s in e.args if isinstance(s, str)]
             if len(msg_list) > 0:
                 for s in msg_list:
-                    if 'Failed to add concrete function' in s \
-                        or "Tried to export a function which references an 'untracked' resource" in s:
-                        export_archive = tf_keras.export.ExportArchive()
-                        export_archive.add_endpoint(
-                            name=SIGNATURE_KEY,
-                            fn=lambda *inputs : model(inputs),
-                            input_signature=[tf.TensorSpec(tensor.shape, tensor.dtype, tensor.name) for tensor in model.inputs],
+                    if 'A root scope name has to match the following pattern' in s:
+                        error(
+                            f'Generation of saved_model failed because the OP name does not match the following pattern. ^[A-Za-z0-9.][A-Za-z0-9_.\\\\/>-]*$'
                         )
-                        export_archive.write_out(output_folder_path)
-                        break
+                        matches = re.findall(r"'([^']*)'", s)
+                        error(f'{matches[0]}')
+                        error(
+                            f'Please convert again with the `-osd` or `--output_signaturedefs` option.'
+                        )
+                        sys.exit(1)
             else:
                 error(e)
                 import traceback
@@ -1408,19 +1526,23 @@ def convert(
                 )
 
         # Quantized TFLite
-        MEAN = np.asarray([[[[0.485, 0.456, 0.406]]]], dtype=np.float32)
-        STD = np.asarray([[[[0.229, 0.224, 0.225]]]], dtype=np.float32)
         if output_integer_quantized_tflite:
             # Get signatures/input keys
-            loaded_saved_model = tf.saved_model.load(
-                output_folder_path
-            ).signatures[SIGNATURE_KEY]
-            input_keys = list(loaded_saved_model.structured_input_signature[1].keys())
-            input_shapes = [v.shape for v in loaded_saved_model.structured_input_signature[1].values()]
-            input_dtypes = [v.dtype for v in loaded_saved_model.structured_input_signature[1].values()]
-            output_keys = list(loaded_saved_model.structured_outputs.keys())
-            output_shapes = [v.shape for v in loaded_saved_model.structured_outputs.values()]
-            output_dtypes = [v.dtype for v in loaded_saved_model.structured_outputs.values()]
+            trackable_obj = \
+                tf.saved_model.load(
+                    output_folder_path
+                )
+            loaded_saved_model: _WrapperFunction = trackable_obj.signatures[SIGNATURE_KEY]
+            structured_input_signature: Dict[str, tf.TensorSpec] = loaded_saved_model.structured_input_signature[1]
+            structured_outputs: Dict[str, tf.TensorSpec] = loaded_saved_model.structured_outputs
+
+            input_keys: List[str] = list(structured_input_signature.keys())
+            input_shapes: List[tf.TensorShape] = [v.shape for v in structured_input_signature.values()]
+            input_dtypes: List[tf.dtypes.DType] = [v.dtype for v in structured_input_signature.values()]
+
+            output_keys: List[str] = list(structured_outputs.keys())
+            output_shapes: List[tf.TensorShape] = [v.shape for v in structured_outputs.values()]
+            output_dtypes: List[tf.dtypes.DType] = [v.dtype for v in structured_outputs.values()]
 
             print('')
             info(Color.BLUE(f'Signature information for quantization'))
@@ -1463,15 +1585,30 @@ def convert(
                 for model_input in model.inputs:
                     if model_input.dtype != tf.float32 \
                         or len(model_input.shape) != 4 \
-                        or model_input.shape[-1] != 3:
+                        or model_input.shape[-1] not in [3, 4]:
                         error(
                             f'For models that have multiple input OPs and need to perform INT8 quantization calibration '+
-                            f'using non-rgb-image input tensors, specify the calibration data with '+
+                            f'using non-rgb-image/non-rgba-image input tensors, specify the calibration data with '+
                             f'--quant_calib_input_op_name_np_data_path. '+
                             f'model_input[n].shape: {model_input.shape}'
                         )
                         sys.exit(1)
 
+                    if model_input.shape[-1] == 3:
+                        # RGB
+                        mean = np.asarray([[[[0.485, 0.456, 0.406]]]], dtype=np.float32)
+                        std = np.asarray([[[[0.229, 0.224, 0.225]]]], dtype=np.float32)
+                    elif model_input.shape[-1] == 4:
+                        # RGBA
+                        mean = np.asarray([[[[0.485, 0.456, 0.406, 0.000]]]], dtype=np.float32)
+                        std = np.asarray([[[[0.229, 0.224, 0.225, 1.000]]]], dtype=np.float32)
+                        new_element_array = np.full((*calib_data.shape[:-1], 1), 0.500, dtype=np.float32)
+                        calib_data = np.concatenate((calib_data, new_element_array), axis=-1)
+                    else:
+                        # Others
+                        mean = np.asarray([[[[0.485, 0.456, 0.406]]]], dtype=np.float32)
+                        std = np.asarray([[[[0.229, 0.224, 0.225]]]], dtype=np.float32)
+
                     calib_data_dict[model_input.name] = \
                         [
                             tf.image.resize(
@@ -1481,8 +1618,8 @@ def convert(
                                     model_input.shape[2] if model_input.shape[2] is not None else 640,
                                 )
                             ),
-                            MEAN,
-                            STD,
+                            mean,
+                            std,
                         ]
             elif custom_input_op_name_np_data_path is not None:
                 for param in custom_input_op_name_np_data_path:
@@ -1846,6 +1983,7 @@ def main():
         '-osd',
         '--output_signaturedefs',
         action='store_true',
+        default=True,
         help=\
             'Signature is added to the output for serving or for conversion \n' +
             'to other model formats. However, this can significantly reduce the speed \n' +
@@ -2101,6 +2239,18 @@ def main():
             'e.g. \n' +
             '--keep_shape_absolutely_input_names "input0" "input1" "input2"'
     )
+    parser.add_argument(
+        '-inimc',
+        '--input_names_to_interrupt_model_conversion',
+        type=str,
+        nargs='+',
+        help=\
+            'Input names that interrupt model conversion. \n' +
+            'Interrupts model transformation at the specified input name \n' +
+            'and inputs the model partitioned into subgraphs. \n\n' +
+            'e.g. \n' +
+            '--input_names_to_interrupt_model_conversion "input0" "input1" "input2"'
+    )
     parser.add_argument(
         '-onimc',
         '--output_names_to_interrupt_model_conversion',
@@ -2444,6 +2594,7 @@ def main():
         keep_ncw_or_nchw_or_ncdhw_input_names=args.keep_ncw_or_nchw_or_ncdhw_input_names,
         keep_nwc_or_nhwc_or_ndhwc_input_names=args.keep_nwc_or_nhwc_or_ndhwc_input_names,
         keep_shape_absolutely_input_names=args.keep_shape_absolutely_input_names,
+        input_names_to_interrupt_model_conversion=args.input_names_to_interrupt_model_conversion,
         output_names_to_interrupt_model_conversion=args.output_names_to_interrupt_model_conversion,
         disable_group_convolution=args.disable_group_convolution,
         enable_accumulation_type_float16=args.enable_accumulation_type_float16,

onnx2tf/ops/Add.py

@@ -168,8 +168,37 @@ def make_node(
         is_scalar_2_rank = tf.rank(input_tensor_2) == 0
         if hasattr(is_scalar_2_rank, 'numpy'):
             is_scalar_2 = is_scalar_2_rank.numpy()
+
         if (is_scalar_1 or is_scalar_2) and graph_node.i().op == 'Gemm':
             pass
+        elif (is_scalar_1 or is_scalar_2) and graph_node.i().op != 'Gemm':
+            first_tensor = None
+            second_tensor = None
+            if is_scalar_1:
+                first_tensor = input_tensor_2
+                second_tensor = input_tensor_1
+            elif is_scalar_2:
+                first_tensor = input_tensor_1
+                second_tensor = input_tensor_2
+            tmp_result = tf.math.add(first_tensor, second_tensor)
+            tmp_result_shape = tmp_result.shape
+            if first_tensor.shape == tmp_result_shape:
+                pass
+            else:
+                input_tensor_1, input_tensor_2 = \
+                    pre_explicit_broadcast(
+                        input_tensor_1=input_tensor_1,
+                        input_tensor_2=input_tensor_2,
+                    )
+
+                input_tensor_1, input_tensor_2 = \
+                    explicit_broadcast(
+                        const_or_var_1=input_tensor_1,
+                        const_or_var_2=input_tensor_2,
+                        graph_node=graph_node,
+                        tf_layers_dict= tf_layers_dict,
+                    )
+
         else:
             input_tensor_1, input_tensor_2 = \
                 pre_explicit_broadcast(

onnx2tf/ops/AveragePool.py

@@ -268,11 +268,14 @@ def make_node(
             [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',
-        )
+        if spatial_size == 1 and kernel_shape[0] > input_tensor_shape[1]:
+            padded_tensor = input_tensor
+        else:
+            padded_tensor = tf.pad(
+                tensor=input_tensor,
+                paddings=tf_pads,
+                mode='CONSTANT',
+            )
 
     else:
         padded_tensor = input_tensor
@@ -306,11 +309,18 @@ def make_node(
     # Generation of TF OP
     tf_op_type = None
     if len(kernel_shape) == 1:
-        pooled_tensor = AveragePooling1D(
-            pool_size=kernel_shape,
-            strides=strides,
-            padding=tf_pad_mode.upper(),
-        )(padded_tensor)
+        if kernel_shape[0] > padded_tensor.shape[1]:
+            pooled_tensor = AveragePooling1D(
+                pool_size=[padded_tensor.shape[1]],
+                strides=[padded_tensor.shape[1]],
+                padding=tf_pad_mode.upper(),
+            )(padded_tensor)
+        else:
+            pooled_tensor = AveragePooling1D(
+                pool_size=kernel_shape,
+                strides=strides,
+                padding=tf_pad_mode.upper(),
+            )(padded_tensor)
         tf_op_type = AveragePooling1D
 
     elif len(kernel_shape) == 2: