Awesome
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
- https://github.com/onnx/onnx/blob/main/docs/Operators.md
- https://docs.nvidia.com/deeplearning/tensorrt/onnx-graphsurgeon/docs/index.html
- https://github.com/NVIDIA/TensorRT/tree/main/tools/onnx-graphsurgeon
- https://github.com/PINTO0309/simple-onnx-processing-tools
- https://github.com/PINTO0309/PINTO_model_zoo
8. Issues
https://github.com/PINTO0309/simple-onnx-processing-tools/issues