From 1a1e2cce4c90b1e70315dcb298d852ff273cee85 Mon Sep 17 00:00:00 2001 From: pinto0309 Date: Wed, 6 Apr 2022 14:20:24 +0900 Subject: [PATCH] Added option to auto downcast --- README.md | 2 +- scs4onnx/__init__.py | 2 +- scs4onnx/onnx_shrink_constant.py | 22 +++++++++++++++++++++- 3 files changed, 23 insertions(+), 3 deletions(-) diff --git a/README.md b/README.md index 8659e25..3cbd7eb 100644 --- a/README.md +++ b/README.md @@ -11,7 +11,7 @@ A very simple tool that compresses the overall size of the ONNX model by aggrega - [ ] ~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. - [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. -- [ ] Added option to downcast from Float64 to Float32 and INT64 to INT32 to attempt size compression. +- [x] Added option to downcast from Float64 to Float32 and INT64 to INT32 to attempt size compression. - [ ] 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. diff --git a/scs4onnx/__init__.py b/scs4onnx/__init__.py index 8a87a50..f22e99a 100644 --- a/scs4onnx/__init__.py +++ b/scs4onnx/__init__.py @@ -1,3 +1,3 @@ from scs4onnx.onnx_shrink_constant import shrinking, main -__version__ = '1.0.8' \ No newline at end of file +__version__ = '1.0.9' \ No newline at end of file diff --git a/scs4onnx/onnx_shrink_constant.py b/scs4onnx/onnx_shrink_constant.py index e1be649..bcff8fe 100644 --- a/scs4onnx/onnx_shrink_constant.py +++ b/scs4onnx/onnx_shrink_constant.py @@ -2,6 +2,7 @@ import os import sys +import traceback from pprint import pprint from argparse import ArgumentParser import numpy as np @@ -112,6 +113,22 @@ def shrinking( continue if np.isscalar(graph_node_input.values): continue + + # Try downcast + ### INT64 -> INT32 + if graph_node_input.values.dtype == np.int64: + orig = graph_node_input.values + dist = graph_node_input.values.astype(np.int32) + if (orig == dist).all(): + graph_node_input.values = dist + + ### Float64 -> Float32 + if graph_node_input.values.dtype == np.float64: + orig = graph_node_input.values + dist = graph_node_input.values.astype(np.float32) + if (orig == dist).all(): + graph_node_input.values = dist + constants[graph_node_input.name] = graph_node_input if not non_verbose: print( @@ -251,7 +268,7 @@ def shrinking( new_model = None try: new_model = onnx.shape_inference.infer_shapes(shrunken_graph) - except: + except Exception as e: new_model = shrunken_graph if not non_verbose: print( @@ -259,6 +276,8 @@ def shrinking( 'The input shape of the next OP does not match the output shape. '+ 'Be sure to open the .onnx file to verify the certainty of the geometry.' ) + tracetxt = traceback.format_exc().splitlines()[-1] + print(f'{Color.YELLOW}WARNING:{Color.RESET} {tracetxt}') # Save if output_onnx_file_path: @@ -299,6 +318,7 @@ def main(): parser.add_argument( '--forced_extraction_op_names', type=str, + default='', help="\ Extracts the constant value of the specified OP name to .npy \ regardless of the mode specified. \