Implementation of a specification for separating the weight of a spec…

…ified OP name to an external file

README.md

@@ -9,7 +9,8 @@ A very simple tool that compresses the overall size of the ONNX model by aggrega
 - [x] Ignore scalar values.
 - [x] Ignore variables.
 - [ ] ~Finally, create a Fork of **[onnx-simplifier](https://github.com/daquexian/onnx-simplifier)** and merge this process just before the onnx file output process~ -> Temporarily abandoned because it turned out that the onnx-simplifier specification needed to be changed in a major way.
-- [ ] Implementation of a specification for separating the weight of a specified OP name to an external file.
+- [x] Implementation of a specification for separating the weight of a specified OP name to an external file.
+- [ ] Implementation of a specification for separating the weight of a specified Constant name to an external file.
 - [ ] Final work-around idea for breaking the 2GB limit, since the internal logic of onnx has a Protocol Buffers limit of 2GB checked. Recombine after optimization. Splitting and merging seems like it would be easy. For each partitioned onnx component, optimization is performed in the order of onnx-simplifier → scs4onnx to optimize the structure while keeping the buffer size to a minimum, and then the optimized components are recombined to reconstruct the whole graph. Finally, run scs4onnx again on the reconstructed, optimized overall graph to further reduce the model-wide constant.
 
 
@@ -43,7 +44,12 @@ $ cd /workdir
 $ scs4onnx -h
 
 usage:
-scs4onnx [-h] [--mode {shrink,npy}] [--non_verbose] input_onnx_file_path output_onnx_file_path
+  scs4onnx [-h]
+  [--mode {shrink,npy}]
+  [--forced_extraction_op_names FORCED_EXTRACTION_OP_NAMES]
+  [--non_verbose]
+  input_onnx_file_path output_onnx_file_path
+
 
 positional arguments:
   input_onnx_file_path
@@ -52,8 +58,10 @@ positional arguments:
                         Output onnx file path.
 
 optional arguments:
-  -h, --help            show this help message and exit
-  --mode {shrink,npy}   Constant Value Compression Mode.
+  -h, --help
+                        show this help message and exit
+  --mode {shrink,npy}
+                        Constant Value Compression Mode.
                         shrink: Share constant values inside the model as much as possible.
                                 The model size is slightly larger because
                                 some shared constant values remain inside the model,
@@ -62,7 +70,13 @@ optional arguments:
                                 external file .npy. Instead of the smallest model body size,
                                 the file loading overhead is greater.
                         Default: shrink
-  --non_verbose         Do not show all information logs. Only error logs are displayed.
+  --forced_extraction_op_names FORCED_EXTRACTION_OP_NAMES
+                        Extracts the constant value of the specified OP name to .npy
+                        regardless of the mode specified.
+                        Specify the name of the OP, separated by commas.
+                        e.g. --forced_extraction_op_names aaa,bbb,ccc
+  --non_verbose
+                        Do not show all information logs. Only error logs are displayed.
 ```
 
 ## 3. In-script Usage
@@ -78,6 +92,7 @@ shrinking(
   output_onnx_file_path: Union[str, NoneType] = '',
   onnx_graph: Union[onnx.onnx_ml_pb2.ModelProto, NoneType] = None,
   mode: Union[str, NoneType] = 'shrink',
+  forced_extraction_op_names: List[str] = [],
   non_verbose: Union[bool, NoneType] = False
 ) -> Tuple[onnx.onnx_ml_pb2.ModelProto, str]
 
@@ -105,6 +120,10 @@ shrinking(
             Instead of the smallest model body size, the file loading overhead is greater.
         Default: shrink
 
+    forced_extraction_op_names: List[str]
+        Extracts the constant value of the specified OP name to .npy
+        regardless of the mode specified. e.g. ['aaa','bbb','ccc']
+
     non_verbose: Optional[bool]
         Do not show all information logs. Only error logs are displayed.
         Default: False

scs4onnx/__init__.py

@@ -1,3 +1,3 @@
 from scs4onnx.onnx_shrink_constant import shrinking, main
 
-__version__ = '1.0.7'
+__version__ = '1.0.8'

scs4onnx/onnx_shrink_constant.py

@@ -41,6 +41,7 @@ def shrinking(
     output_onnx_file_path: Optional[str] = '',
     onnx_graph: Optional[onnx.ModelProto] = None,
     mode: Optional[str] = 'shrink',
+    forced_extraction_op_names: List[str] = [],
     non_verbose: Optional[bool] = False,
 ) -> Tuple[onnx.ModelProto, List[str]]:
 
@@ -69,6 +70,10 @@ def shrinking(
             Instead of the smallest model body size, the file loading overhead is greater.\n\
         Default: shrink
 
+    forced_extraction_op_names: List[str]
+        Extracts the constant value of the specified OP name to .npy regardless of the mode specified.\n\
+        e.g. ['aaa','bbb','ccc']
+
     non_verbose: Optional[bool]
         Do not show all information logs. Only error logs are displayed.\n\
         Default: False
@@ -116,13 +121,13 @@ def shrinking(
 
     constants_list = list(constants.items())
 
+    # Aggregate list generation
+    aggregate_constants_name_list = {}
+    aggregate_constants_value_list = {}
+
     # Processing is performed only when the number of constant tensors to be processed is 2 or more
     if len(constants_list) >= 2:
-        # Aggregate list generation
-        aggregate_constants_name_list = {}
-        aggregate_constants_value_list = {}
         aggregate_constants_matched_list = [False] * len(constants_list)
-
         for comparator_idx in range(0, len(constants_list)-1):
             if not aggregate_constants_matched_list[comparator_idx]:
                 for comparison_dest_idx in range(comparator_idx+1, len(constants_list)):
@@ -142,6 +147,7 @@ def shrinking(
     aggregate_constants
     {'425': ['434', '5506'], '1646': ['1910', '2608', '2872', '3570', '3834'], '5550': ['4107']}
     """
+    # Automatic aggregation of constants
     npy_file_paths = []
     for layer_name_quoted_src, layer_name_quoted_dists in aggregate_constants_name_list.items():
         i = None
@@ -178,6 +184,61 @@ def shrinking(
                         graph.nodes[graph_node_idx].inputs[input_idx] = i
 
     graph.cleanup().toposort()
+
+
+    # Searches the entire model by the OP name specified in forced_extraction_op_names,
+    # and if an OP with a matching name is found, the constant is forced to be extracted to the .npy file.
+
+    # Constant Value Extraction
+    constants = {}
+    graph_nodes = [node for node in graph.nodes if node.name in forced_extraction_op_names]
+    for graph_node in graph_nodes:
+        for graph_node_input in graph_node.inputs:
+            if not isinstance(graph_node_input, Constant):
+                continue
+            if len(graph_node_input.shape) == 0:
+                continue
+            if np.isscalar(graph_node_input.values):
+                continue
+            constants[graph_node_input.name] = graph_node_input
+    if not non_verbose:
+        print(
+            f'{Color.GREEN}INFO:{Color.RESET} '+
+            f'Forced-extraction number of constant values to be studied: {len(constants)}'
+        )
+
+    constants_list = list(constants.items())
+
+    for constant_key, constant_value in constants_list:
+        i = None
+        # Export constant values to a numpy file
+        external_file_name = ''
+        if output_onnx_file_path:
+            external_file_name = \
+                f"{os.path.splitext(os.path.basename(output_onnx_file_path))[0]}" + \
+                f"_exported_{constant_key.replace(':','_').replace(';','_').replace('/','_').replace(',','_')}.npy"
+        else:
+            external_file_name = \
+                f"exported_{constant_key.replace(':','_').replace(';','_').replace('/','_').replace(',','_')}.npy"
+        np.save(
+            external_file_name,
+            constant_value.values
+        )
+        npy_file_paths.append(external_file_name)
+        # Generate Inputs
+        i = gs.Variable(
+            name=external_file_name,
+            dtype=constant_value.values.dtype,
+            shape=constant_value.values.shape,
+        )
+        for graph_node_idx, graph_node in enumerate(graph.nodes):
+            for input_idx, graph_node_input in enumerate(graph_node.inputs):
+                if graph_node_input.name == constant_key:
+                    graph.nodes[graph_node_idx].inputs[input_idx] = i
+        graph.inputs.append(i)
+
+    graph.cleanup().toposort()
+
     if not non_verbose:
         print(f'{Color.GREEN}INFO:{Color.RESET} Results:')
         pprint(aggregate_constants_name_list)
@@ -235,6 +296,15 @@ def main():
             Instead of the smallest model body size, the file loading overhead is greater. \
             Default: shrink"
     )
+    parser.add_argument(
+        '--forced_extraction_op_names',
+        type=str,
+        help="\
+            Extracts the constant value of the specified OP name to .npy \
+            regardless of the mode specified. \
+            Specify the name of the OP, separated by commas. \
+            e.g. --forced_extraction_op_names aaa,bbb,ccc"
+    )
     parser.add_argument(
         '--non_verbose',
         action='store_true',
@@ -253,13 +323,19 @@ def main():
         )
         sys.exit(1)
 
+    forced_extraction_op_names = args.forced_extraction_op_names.strip(' ,').replace(' ','').split(',')
+
     # Model shrink
     shrunken_graph, npy_file_paths = shrinking(
         input_onnx_file_path=args.input_onnx_file_path,
         output_onnx_file_path=args.output_onnx_file_path,
         mode=args.mode,
+        forced_extraction_op_names=forced_extraction_op_names,
         non_verbose=args.non_verbose
     )
 
     if not args.non_verbose:
         print(f'{Color.GREEN}INFO:{Color.RESET} Finish!')
+
+if __name__ == '__main__':
+    main()