sna4onnx

Simple node addition tool for onnx. Simple Node Addition for ONNX.

https://github.com/PINTO0309/simple-onnx-processing-tools

Key concept

• Combines the OP generated by sog4onnx with the specified output and input OPs.
• Use a good combination with the ONNX merge tool snc4onnx while merging OP parts.
• Only one OP can be extrapolated at a time.
• After OP extrapolation, the entire model is checked, and even if there is a consistency problem, only a warning is displayed and the ONNX file is output as is.
• Add unconnected input and output variables to the input/output OP of a graph.

1. Setup

1-1. HostPC

### option
$ echo export PATH="~/.local/bin:$PATH" >> ~/.bashrc \
&& source ~/.bashrc

### run
$ pip install -U onnx \
&& python3 -m pip install -U onnx_graphsurgeon --index-url https://pypi.ngc.nvidia.com \
&& pip install -U sna4onnx

1-2. Docker

https://github.com/PINTO0309/simple-onnx-processing-tools#docker

2. CLI Usage

$ sna4onnx -h

usage:
  sna4onnx [-h]
  -if INPUT_ONNX_FILE_PATH
  -aot ADD_OP_TYPE
  -aon ADD_OP_NAME
  [-aoiv NAME TYPE VALUE]
  [-aoov NAME TYPE VALUE]
  [-aoa NAME DTYPE VALUE]
  [-csoon SRCOP_NAME SRCOP_OUTPUT_NAME ADDOP_NAME ADDOP_INPUT_NAME]
  -cdoin ADDOP_NAME ADDOP_OUTPUT_NAME DESTOP_NAME DESTOP_INPUT_NAME
  [-of OUTPUT_ONNX_FILE_PATH]
  [-n]

optional arguments:
  -h, --help
      show this help message and exit

  -if INPUT_ONNX_FILE_PATH, --input_onnx_file_path INPUT_ONNX_FILE_PATH
      Input onnx file path.

  -aot ADD_OP_TYPE, --add_op_type ADD_OP_TYPE
      ONNX OP type.
      https://github.com/onnx/onnx/blob/main/docs/Operators.md

  -aon ADD_OP_NAME, --add_op_name ADD_OP_NAME
      Name of OP to be added.
      e.g.
      --add_op_name AddOP1

  -aoiv ADD_OP_INPUT_VARIABLES ADD_OP_INPUT_VARIABLES ADD_OP_INPUT_VARIABLES, --add_op_input_variables ADD_OP_INPUT_VARIABLES ADD_OP_INPUT_VARIABLES ADD_OP_INPUT_VARIABLES
      input_variables can be specified multiple times.
      --add_op_input_variables variable_name numpy.dtype shape
      https://github.com/onnx/onnx/blob/main/docs/Operators.md
      e.g.
      --add_op_input_variables inpname1 float32 [1,3,5,5]
      --add_op_input_variables inpname2 int32 [1]
      --add_op_input_variables inpname3 float64 [1,3,224,224]

  -aoov ADD_OP_OUTPUT_VARIABLES ADD_OP_OUTPUT_VARIABLES ADD_OP_OUTPUT_VARIABLES, --add_op_output_variables ADD_OP_OUTPUT_VARIABLES ADD_OP_OUTPUT_VARIABLES ADD_OP_OUTPUT_VARIABLES
      output_variables can be specified multiple times.
      --add_op_output_variables variable_name numpy.dtype shape
      https://github.com/onnx/onnx/blob/main/docs/Operators.md
      e.g.
      --add_op_output_variables outname1 float32 [1,3,5,5]
      --add_op_output_variables outname2 int32 [1]
      --add_op_output_variables outname3 float64 [1,3,224,224]

  -aoa ADD_OP_ATTRIBUTES ADD_OP_ATTRIBUTES ADD_OP_ATTRIBUTES, --add_op_attributes ADD_OP_ATTRIBUTES ADD_OP_ATTRIBUTES ADD_OP_ATTRIBUTES
      attributes can be specified multiple times.
      --add_op_attributes name dtype value
      dtype is one of "float32" or "float64" or "int32" or "int64" or "str".
      https://github.com/onnx/onnx/blob/main/docs/Operators.md
      e.g.
      --add_op_attributes alpha float32 1.0
      --add_op_attributes beta float32 1.0
      --add_op_attributes transA int64 0
      --add_op_attributes transB int64 0

  -csoon CONNECTION_SRC_OP_OUTPUT_NAMES CONNECTION_SRC_OP_OUTPUT_NAMES CONNECTION_SRC_OP_OUTPUT_NAMES CONNECTION_SRC_OP_OUTPUT_NAMES, --connection_src_op_output_names CONNECTION_SRC_OP_OUTPUT_NAMES CONNECTION_SRC_OP_OUTPUT_NAMES CONNECTION_SRC_OP_OUTPUT_NAMES CONNECTION_SRC_OP_OUTPUT_NAMES
      Specify the name of the output name from which to connect.
      e.g.
      -Before-
        [OpA]oname1 - iname1[OpB]oname1
        [OpC]oname1

      -After-
        [OpA]oname1 - iname1[AddOP1]oname1 - iname1[OpB]oname1
        [OpC]oname1 - iname2[AddOP1]

      When extrapolating a new OP between OpA and OpB.
      --connection_src_op_output_names OpA oname1 AddOP1 iname1
      --connection_src_op_output_names OpC oname1 AddOP1 iname2
      This need not be specified only when the type of the newly added OP is Constant.

  -cdoin CONNECTION_DEST_OP_INPUT_NAMES CONNECTION_DEST_OP_INPUT_NAMES CONNECTION_DEST_OP_INPUT_NAMES CONNECTION_DEST_OP_INPUT_NAMES, --connection_dest_op_input_names CONNECTION_DEST_OP_INPUT_NAMES CONNECTION_DEST_OP_INPUT_NAMES CONNECTION_DEST_OP_INPUT_NAMES CONNECTION_DEST_OP_INPUT_NAMES
      Specify the name of the input name from which to connect.
      e.g.
      -Before-
        [OpA]oname1 - iname1[OpB]oname1
        [OpC]oname1

      -After-
        [OpA]oname1 - iname1[AddOP1]oname1 - iname1[OpB]oname1
        [OpC]oname1 - iname2[AddOP1]

      When extrapolating a new OP between OpA and OpB.
      --connection_dest_op_input_names AddOP1 oname1 OpB iname1

  -of OUTPUT_ONNX_FILE_PATH, --output_onnx_file_path OUTPUT_ONNX_FILE_PATH
      Output onnx file path.
      If not specified, a file with "_mod" appended to the end of input_onnx_file_path is output.
      e.g.
      aaa.onnx -> aaa_mod.onnx

  -n, --non_verbose
      Do not show all information logs. Only error logs are displayed.

3. In-script Usage

>>> from sna4onnx import add
>>> help(add)

Help on function add in module sna4onnx.onnx_operation_adder:

add(
  connection_src_op_output_names: List,
  connection_dest_op_input_names: List,
  add_op_type: str, add_op_name: str,
  add_op_input_variables: Union[dict, NoneType] = None,
  add_op_output_variables: Union[dict, NoneType] = None,
  add_op_attributes: Union[dict, NoneType] = None,
  input_onnx_file_path: Union[str, NoneType] = '',
  onnx_graph: Union[onnx.onnx_ml_pb2.ModelProto, NoneType] = None,
  output_onnx_file_path: Union[str, NoneType] = '',
  non_verbose: Union[bool, NoneType] = False
) -> onnx.onnx_ml_pb2.ModelProto

    Parameters
    ----------
    connection_src_op_output_names: List
        Specify the name of the output name from which to connect.

        e.g.
        -Before-
            [OpA] outnameA - inpnameB1 [OpB] outnameB
            [OpC] outnameC
        -After-
            [OpA] outnameA - inpname1 [AddOP1] outname1 - inpnameB1 [OpB] outnameB
            [OpC] outnameC - inpname2 [AddOP1]
        When extrapolating a new OP between OpA and OpB.
        connection_src_op_output_names = [
            ["OpA", "outnameA", "AddOP1", "inpname1",],
            ["OpC", "outnameC", "AddOP1", "inpname2",],
        ]

        This need not be specified only when the type of the newly added OP is Constant.

    connection_dest_op_input_names: List
        Specify the name of the input name from which to connect.

        e.g.
        -Before-
            [OpA] outnameA - inpnameB1 [OpB] outnameB
            [OpC] outnameC
        -After-
            [OpA] outnameA - inpname1 [AddOP1] outname1 - inpnameB1 [OpB] outnameB
            [OpC] outnameC - inpname2 [AddOP1]
        When extrapolating a new OP between OpA and OpB.
        connection_dest_op_input_names = [
            ["AddOP1", "outname1", "OpB", "inpnameB1"],
        ]

    add_op_type: str
        ONNX op type.
        See below for the types of OPs that can be specified.

        e.g. "Add", "Div", "Gemm", ...
        https://github.com/onnx/onnx/blob/main/docs/Operators.md

    add_op_name: str
        Name of OP to be added.

        e.g. --add_op_name AddOP1

    add_op_input_variables: Optional[dict]
        Specify input variables for the OP to be generated.
        See below for the variables that can be specified.

        {
            "input_var_name1": [numpy.dtype, shape],
            "input_var_name2": [numpy.dtype, shape],
            ...
        }

        e.g.
        add_op_input_variables = {
            "inpname1": [np.float32, [1,224,224,3]],
            "inpname2": [np.bool_, [0]],
            ...
        }
        https://github.com/onnx/onnx/blob/main/docs/Operators.md

    add_op_output_variables: Optional[dict]
        Specify output variables for the OP to be generated.
        See below for the variables that can be specified.

        {
            "output_var_name1": [numpy.dtype, shape],
            "output_var_name2": [numpy.dtype, shape],
            ...
        }

        e.g.
        add_op_output_variables = {
            "outname1": [np.float32, [1,224,224,3]],
            "outname2": [np.bool_, [0]],
            ...
        }
        https://github.com/onnx/onnx/blob/main/docs/Operators.md

    add_op_attributes: Optional[dict]
        Specify output add_op_attributes for the OP to be generated.
        See below for the add_op_attributes that can be specified.

        {
            "attr_name1": value1,
            "attr_name2": value2,
            "attr_name3": value3,
            ...
        }

        e.g.
        add_op_attributes = {
            "alpha": 1.0,
            "beta": 1.0,
            "transA": 0,
            "transB": 0,
        }
        Default: None
        https://github.com/onnx/onnx/blob/main/docs/Operators.md

    input_onnx_file_path: Optional[str]
        Input onnx file path.
        Either input_onnx_file_path or onnx_graph must be specified.
        Default: ''

    onnx_graph: Optional[onnx.ModelProto]
        onnx.ModelProto.
        Either input_onnx_file_path or onnx_graph must be specified.
        onnx_graph If specified, ignore input_onnx_file_path and process onnx_graph.

    output_onnx_file_path: Optional[str]
        Output onnx file path. If not specified, no ONNX file is output.
        Default: ''

    non_verbose: Optional[bool]
        Do not show all information logs. Only error logs are displayed.
        Default: False

    Returns
    -------
    changed_graph: onnx.ModelProto
        Changed onnx ModelProto.

4. CLI Execution

$ sna4onnx \
--input_onnx_file_path crestereo_init_iter2_120x160.onnx \
--connection_src_op_output_name Sub_451 onnx::Pow_603 dummy_mul inp1 \
--connection_dest_op_input_name dummy_mul out1 Div_458 onnx::Pow_603 \
--add_op_type Mul \
--add_op_name dummy_mul \
--add_op_input_variables inp1 float32 [1,70,256] \
--add_op_input_variables inp2_const float32 [1] \
--add_op_output_variables out1 float32 [1,70,256] \
--output_onnx_file_path test.onnx

5. In-script Execution

from sna4onnx import add

onnx_graph = add(
  input_onnx_file_path="crestereo_init_iter2_120x160.onnx",
  connection_src_op_output_names=[
    ["Sub_451", "onnx::Pow_603", "dummy_mul", "inp1"],
  ],
  connection_dest_op_input_names=[
    ["dummy_mul","out1", "Div_458", "onnx::Pow_603"],
  ],
  add_op_type="Mul",
  add_op_name="dummy_mul",
  add_op_input_variables={
    "inp1": [np.float32, [1,70,256]],
    "inp2_const": [np.float32, [1]],
  },
  add_op_output_variables={
    "out1": [np.float32, [1,70,256]],
  },
)

# or

onnx_graph = add(
  onnx_graph=graph,
  connection_src_op_output_names=[
    ["Sub_451", "onnx::Pow_603", "dummy_mul", "inp1"],
  ],
  connection_dest_op_input_names=[
    ["dummy_mul","out1", "Div_458", "onnx::Pow_603"],
  ],
  add_op_type="Mul",
  add_op_name="dummy_mul",
  add_op_input_variables={
    "inp1": [np.float32, [1,70,256]],
    "inp2_const": [np.float32, [1]],
  },
  add_op_output_variables={
    "out1": [np.float32, [1,70,256]],
  },
)

6. Sample

Before

After

7. Reference

1. https://github.com/onnx/onnx/blob/main/docs/Operators.md
2. https://docs.nvidia.com/deeplearning/tensorrt/onnx-graphsurgeon/docs/index.html
3. https://github.com/NVIDIA/TensorRT/tree/main/tools/onnx-graphsurgeon
4. https://github.com/PINTO0309/simple-onnx-processing-tools
5. https://github.com/PINTO0309/PINTO_model_zoo

8. Issues

https://github.com/PINTO0309/simple-onnx-processing-tools/issues