Fine-tuning of document and logic indentation

Author: PINTO0309

README.md

@@ -7,7 +7,7 @@ https://github.com/PINTO0309/simple-onnx-processing-tools
 
 # Key concept
 - [x] If INPUT OP name and OUTPUT OP name are specified, the onnx graph within the range of the specified OP name is extracted and .onnx is generated.
-- [x] Change backend to `onnx.utils.Extractor.extract_model` so that onnx.ModelProto can be specified as input.
+- [x] I do not use `onnx.utils.extractor.extract_model` because it is very slow and I implement my own model separation logic.
 
 ## 1. Setup
 ### 1-1. HostPC
@@ -18,6 +18,7 @@ $ echo export PATH="~/.local/bin:$PATH" >> ~/.bashrc \
 
 ### run
 $ pip install -U onnx \
+&& python3 -m pip install -U onnx_graphsurgeon --index-url https://pypi.ngc.nvidia.com
 && pip install -U sne4onnx
 ```
 ### 1-2. Docker
@@ -37,18 +38,18 @@ usage:
 
 optional arguments:
   -h, --help
-        show this help message and exit
+        show this help message and exit.
 
   --input_onnx_file_path INPUT_ONNX_FILE_PATH
         Input onnx file path.
 
   --input_op_names INPUT_OP_NAMES
         List of OP names to specify for the input layer of the model.
-         e.g. --input_op_names aaa bbb ccc
+        e.g. --input_op_names aaa bbb ccc
 
   --output_op_names OUTPUT_OP_NAMES
         List of OP names to specify for the output layer of the model.
-         e.g. --output_op_names ddd eee fff
+        e.g. --output_op_names ddd eee fff
 
   --output_onnx_file_path OUTPUT_ONNX_FILE_PATH
         Output onnx file path. If not specified, extracted.onnx is output.
@@ -124,8 +125,8 @@ $ sne4onnx \
 from sne4onnx import extraction
 
 extracted_graph = extraction(
-  input_op_names=['aaa', 'bbb', 'ccc'],
-  output_op_names=['ddd', 'eee', 'fff'],
+  input_op_names=['aaa','bbb','ccc'],
+  output_op_names=['ddd','eee','fff'],
   input_onnx_file_path='input.onnx',
   output_onnx_file_path='output.onnx',
 )
@@ -135,8 +136,8 @@ extracted_graph = extraction(
 from sne4onnx import extraction
 
 extracted_graph = extraction(
-  input_op_names=['aaa', 'bbb', 'ccc'],
-  output_op_names=['ddd', 'eee', 'fff'],
+  input_op_names=['aaa','bbb','ccc'],
+  output_op_names=['ddd','eee','fff'],
   onnx_graph=graph,
   output_onnx_file_path='output.onnx',
 )

sne4onnx/__init__.py

@@ -1,3 +1,3 @@
 from sne4onnx.onnx_network_extraction import extraction, main
 
-__version__ = '1.0.6'
+__version__ = '1.0.7'

sne4onnx/onnx_network_extraction.py

@@ -107,16 +107,35 @@ def extraction(
     graph.cleanup().toposort()
 
     # Extraction of input OP and output OP
-    graph_node_inputs = [graph_nodes for graph_nodes in graph.nodes for graph_nodes_input in graph_nodes.inputs if graph_nodes_input.name in input_op_names]
-    graph_node_outputs = [graph_nodes for graph_nodes in graph.nodes for graph_nodes_output in graph_nodes.outputs if graph_nodes_output.name in output_op_names]
+    graph_node_inputs = [
+        graph_nodes \
+            for graph_nodes in graph.nodes \
+                for graph_nodes_input in graph_nodes.inputs \
+                    if graph_nodes_input.name in input_op_names
+    ]
+    graph_node_outputs = [
+        graph_nodes \
+            for graph_nodes in graph.nodes \
+                for graph_nodes_output in graph_nodes.outputs \
+                    if graph_nodes_output.name in output_op_names
+    ]
 
     # Init graph INPUT/OUTPUT
     graph.inputs.clear()
     graph.outputs.clear()
 
     # Update graph INPUT/OUTPUT
-    graph.inputs = [graph_node_input for graph_node in graph_node_inputs for graph_node_input in graph_node.inputs if graph_node_input.shape]
-    graph.outputs = [graph_node_output for graph_node in graph_node_outputs for graph_node_output in graph_node.outputs]
+    graph.inputs = [
+        graph_node_input \
+            for graph_node in graph_node_inputs \
+                for graph_node_input in graph_node.inputs \
+                    if graph_node_input.shape
+    ]
+    graph.outputs = [
+        graph_node_output \
+            for graph_node in graph_node_outputs \
+                for graph_node_output in graph_node.outputs
+    ]
 
     # Cleanup
     graph.cleanup().toposort()