Skip to content

Commit 16cb20c

Browse files
committed
Support for onnx.ModelProto input, optional bug fixes
1 parent 7d7be5b commit 16cb20c

File tree

3 files changed

+26
-22
lines changed

3 files changed

+26
-22
lines changed

README.md

Lines changed: 14 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -72,18 +72,15 @@ $ python
7272
Help on function extraction in module sne4onnx.onnx_network_extraction:
7373

7474
extraction(
75-
input_onnx_file_path: str,
7675
input_op_names: List[str],
7776
output_op_names: List[str],
78-
output_onnx_file_path: Union[str, NoneType] = '',
79-
onnx_graph: Union[onnx.onnx_ml_pb2.ModelProto, NoneType] = None
77+
input_onnx_file_path: Union[str, NoneType] = '',
78+
onnx_graph: Union[onnx.onnx_ml_pb2.ModelProto, NoneType] = None,
79+
output_onnx_file_path: Union[str, NoneType] = ''
8080
) -> onnx.onnx_ml_pb2.ModelProto
8181

8282
Parameters
8383
----------
84-
input_onnx_file_path: str
85-
Input onnx file path.
86-
8784
input_op_names: List[str]
8885
List of OP names to specify for the input layer of the model.
8986
Specify the name of the OP, separated by commas.
@@ -94,16 +91,21 @@ extraction(
9491
Specify the name of the OP, separated by commas.
9592
e.g. ['ddd','eee','fff']
9693

97-
output_onnx_file_path: Optional[str]
98-
Output onnx file path.
99-
If not specified, .onnx is not output.
100-
Default: ''
94+
input_onnx_file_path: Optional[str]
95+
Input onnx file path.
96+
Either input_onnx_file_path or onnx_graph must be specified.
97+
onnx_graph If specified, ignore input_onnx_file_path and process onnx_graph.
10198

10299
onnx_graph: Optional[onnx.ModelProto]
103100
onnx.ModelProto.
104101
Either input_onnx_file_path or onnx_graph must be specified.
105102
onnx_graph If specified, ignore input_onnx_file_path and process onnx_graph.
106103

104+
output_onnx_file_path: Optional[str]
105+
Output onnx file path.
106+
If not specified, .onnx is not output.
107+
Default: ''
108+
107109
Returns
108110
-------
109111
extracted_graph: onnx.ModelProto
@@ -125,9 +127,9 @@ $ sne4onnx \
125127
from sne4onnx import extraction
126128

127129
extracted_graph = extraction(
128-
input_onnx_file_path='input.onnx',
129130
input_op_names=['aaa', 'bbb', 'ccc'],
130131
output_op_names=['ddd', 'eee', 'fff'],
132+
input_onnx_file_path='input.onnx',
131133
output_onnx_file_path='output.onnx',
132134
)
133135
```
@@ -138,8 +140,8 @@ from sne4onnx import extraction
138140
extracted_graph = extraction(
139141
input_op_names=['aaa', 'bbb', 'ccc'],
140142
output_op_names=['ddd', 'eee', 'fff'],
141-
output_onnx_file_path='output.onnx',
142143
onnx_graph=graph,
144+
output_onnx_file_path='output.onnx',
143145
)
144146
```
145147

sne4onnx/__init__.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,3 @@
11
from sne4onnx.onnx_network_extraction import extraction, main
22

3-
__version__ = '1.0.4'
3+
__version__ = '1.0.5'

sne4onnx/onnx_network_extraction.py

Lines changed: 11 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -32,19 +32,16 @@ class Color:
3232

3333

3434
def extraction(
35-
input_onnx_file_path: str,
3635
input_op_names: List[str],
3736
output_op_names: List[str],
38-
output_onnx_file_path: Optional[str] = '',
37+
input_onnx_file_path: Optional[str] = '',
3938
onnx_graph: Optional[onnx.ModelProto] = None,
39+
output_onnx_file_path: Optional[str] = '',
4040
) -> onnx.ModelProto:
4141

4242
"""
4343
Parameters
4444
----------
45-
input_onnx_file_path: str
46-
Input onnx file path.
47-
4845
input_op_names: List[str]
4946
List of OP names to specify for the input layer of the model.\n\
5047
Specify the name of the OP, separated by commas.\n\
@@ -55,16 +52,21 @@ def extraction(
5552
Specify the name of the OP, separated by commas.\n\
5653
e.g. ['ddd','eee','fff']
5754
58-
output_onnx_file_path: Optional[str]
59-
Output onnx file path.\n\
60-
If not specified, .onnx is not output.\n\
61-
Default: ''
55+
input_onnx_file_path: Optional[str]
56+
Input onnx file path.\n\
57+
Either input_onnx_file_path or onnx_graph must be specified.\n\
58+
onnx_graph If specified, ignore input_onnx_file_path and process onnx_graph.
6259
6360
onnx_graph: Optional[onnx.ModelProto]
6461
onnx.ModelProto.\n\
6562
Either input_onnx_file_path or onnx_graph must be specified.\n\
6663
onnx_graph If specified, ignore input_onnx_file_path and process onnx_graph.
6764
65+
output_onnx_file_path: Optional[str]
66+
Output onnx file path.\n\
67+
If not specified, .onnx is not output.\n\
68+
Default: ''
69+
6870
Returns
6971
-------
7072
extracted_graph: onnx.ModelProto

0 commit comments

Comments
 (0)