├── requirements.txt
├── logger.py
├── LICENSE
├── README.md
├── .gitignore
└── convert.py


/requirements.txt:
--------------------------------------------------------------------------------
1 | colorlog==4.7.2
2 | onnx==1.8.1
3 | numpy==1.18.5
4 | typer==0.3.2
5 | 


--------------------------------------------------------------------------------
/logger.py:
--------------------------------------------------------------------------------
 1 | import logging
 2 | 
 3 | from colorlog import ColoredFormatter
 4 | 
 5 | LOG_LEVEL = logging.INFO
 6 | LOGFORMAT = "  %(log_color)s%(levelname)-8s%(reset)s| %(name)s | %(log_color)s%(message)s%(reset)s | (%(filename)s:%(lineno)d)"
 7 | logging.root.setLevel(LOG_LEVEL)
 8 | formatter = ColoredFormatter(LOGFORMAT)
 9 | stream = logging.StreamHandler()
10 | stream.setLevel(LOG_LEVEL)
11 | stream.setFormatter(formatter)
12 | log = logging.getLogger("onnx-typecast")
13 | log.setLevel(LOG_LEVEL)
14 | log.addHandler(stream)
15 | 


--------------------------------------------------------------------------------
/LICENSE:
--------------------------------------------------------------------------------
 1 | MIT License
 2 | 
 3 | Copyright (c) 2021 Aadhithya Sankar (ஆதித்யா சங்கர்)
 4 | 
 5 | Permission is hereby granted, free of charge, to any person obtaining a copy
 6 | of this software and associated documentation files (the "Software"), to deal
 7 | in the Software without restriction, including without limitation the rights
 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
 9 | copies of the Software, and to permit persons to whom the Software is
10 | furnished to do so, subject to the following conditions:
11 | 
12 | The above copyright notice and this permission notice shall be included in all
13 | copies or substantial portions of the Software.
14 | 
15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
21 | SOFTWARE.
22 | 


--------------------------------------------------------------------------------
/README.md:
--------------------------------------------------------------------------------
 1 | # onnx-typecast
 2 | A simple python script to typecast ONNX model parameters from INT64 to INT32.
 3 | 
 4 | ## Why?
 5 | I wanted to play around with [ONNX.js](https://github.com/microsoft/onnxjs) and soon figured out that it doesn't support onnx models with INT64 parameters. Also, OpenCV doesn't seem to support INT64 parameters.
 6 | 
 7 | ## What does this script do?
 8 | - The script goes through the parameters of each node of the onnx model and blindly converts INT64 parameters to INT32.
 9 | - It also converts the constant parameters from INT64 to INT32.
10 | - Finally, creates a new model and saves it.
11 | 
12 | ## What's the catch?
13 | - **The script does not handle overflows and blindly converts all INT64 parameters to INT32.**
14 | - So ops that require `>INT32.max` or `<INT32.min` values may not perform as expected.
15 | - Please feel free to modify the script to account for this.
16 | 
17 | ## Alright. How do I use it?
18 |  - simple.
19 |  - Install requirements: `pip install -r requirements.txt`
20 |  - run: `python convert.py path/to/int64_model.onnx path/to/converted_int32_model.onnx`
21 | 
22 | 
23 | ## Also Checkout
24 | - https://github.com/microsoft/onnxjs/issues/168#issuecomment-727219050
25 | 


--------------------------------------------------------------------------------
/.gitignore:
--------------------------------------------------------------------------------
  1 | # Byte-compiled / optimized / DLL files
  2 | __pycache__/
  3 | *.py[cod]
  4 | *$py.class
  5 | 
  6 | # C extensions
  7 | *.so
  8 | 
  9 | # Distribution / packaging
 10 | .Python
 11 | build/
 12 | develop-eggs/
 13 | dist/
 14 | downloads/
 15 | eggs/
 16 | .eggs/
 17 | lib/
 18 | lib64/
 19 | parts/
 20 | sdist/
 21 | var/
 22 | wheels/
 23 | pip-wheel-metadata/
 24 | share/python-wheels/
 25 | *.egg-info/
 26 | .installed.cfg
 27 | *.egg
 28 | MANIFEST
 29 | 
 30 | # PyInstaller
 31 | #  Usually these files are written by a python script from a template
 32 | #  before PyInstaller builds the exe, so as to inject date/other infos into it.
 33 | *.manifest
 34 | *.spec
 35 | 
 36 | # Installer logs
 37 | pip-log.txt
 38 | pip-delete-this-directory.txt
 39 | 
 40 | # Unit test / coverage reports
 41 | htmlcov/
 42 | .tox/
 43 | .nox/
 44 | .coverage
 45 | .coverage.*
 46 | .cache
 47 | nosetests.xml
 48 | coverage.xml
 49 | *.cover
 50 | *.py,cover
 51 | .hypothesis/
 52 | .pytest_cache/
 53 | 
 54 | # Translations
 55 | *.mo
 56 | *.pot
 57 | 
 58 | # Django stuff:
 59 | *.log
 60 | local_settings.py
 61 | db.sqlite3
 62 | db.sqlite3-journal
 63 | 
 64 | # Flask stuff:
 65 | instance/
 66 | .webassets-cache
 67 | 
 68 | # Scrapy stuff:
 69 | .scrapy
 70 | 
 71 | # Sphinx documentation
 72 | docs/_build/
 73 | 
 74 | # PyBuilder
 75 | target/
 76 | 
 77 | # Jupyter Notebook
 78 | .ipynb_checkpoints
 79 | 
 80 | # IPython
 81 | profile_default/
 82 | ipython_config.py
 83 | 
 84 | # pyenv
 85 | .python-version
 86 | 
 87 | # pipenv
 88 | #   According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
 89 | #   However, in case of collaboration, if having platform-specific dependencies or dependencies
 90 | #   having no cross-platform support, pipenv may install dependencies that don't work, or not
 91 | #   install all needed dependencies.
 92 | #Pipfile.lock
 93 | 
 94 | # PEP 582; used by e.g. github.com/David-OConnor/pyflow
 95 | __pypackages__/
 96 | 
 97 | # Celery stuff
 98 | celerybeat-schedule
 99 | celerybeat.pid
100 | 
101 | # SageMath parsed files
102 | *.sage.py
103 | 
104 | # Environments
105 | .env
106 | .venv
107 | env/
108 | venv/
109 | ENV/
110 | env.bak/
111 | venv.bak/
112 | 
113 | # Spyder project settings
114 | .spyderproject
115 | .spyproject
116 | 
117 | # Rope project settings
118 | .ropeproject
119 | 
120 | # mkdocs documentation
121 | /site
122 | 
123 | # mypy
124 | .mypy_cache/
125 | .dmypy.json
126 | dmypy.json
127 | 
128 | # Pyre type checker
129 | .pyre/
130 | 
131 | # custom
132 | onnx_models/
133 | notebooks/
134 | DS_Store/
135 | 


--------------------------------------------------------------------------------
/convert.py:
--------------------------------------------------------------------------------
  1 | import onnx
  2 | 
  3 | from onnx import helper as h
  4 | from onnx import checker as ch
  5 | from onnx import TensorProto, GraphProto
  6 | from onnx import numpy_helper as nph
  7 | 
  8 | import numpy as np
  9 | from collections import OrderedDict
 10 | 
 11 | from logger import log
 12 | import typer
 13 | 
 14 | 
 15 | def make_param_dictionary(initializer):
 16 |     params = OrderedDict()
 17 |     for data in initializer:
 18 |         params[data.name] = data
 19 |     return params
 20 | 
 21 | 
 22 | def convert_params_to_int32(params_dict):
 23 |     converted_params = []
 24 |     for param in params_dict:
 25 |         data = params_dict[param]
 26 |         if data.data_type == TensorProto.INT64:
 27 |             data_cvt = nph.to_array(data).astype(np.int32)
 28 |             data = nph.from_array(data_cvt, data.name)
 29 |         converted_params += [data]
 30 |     return converted_params
 31 | 
 32 | 
 33 | def convert_constant_nodes_to_int32(nodes):
 34 |     """
 35 |     convert_constant_nodes_to_int32 Convert Constant nodes to INT32. If a constant node has data type INT64, a new version of the
 36 |     node is created with INT32 data type and stored.
 37 | 
 38 |     Args:
 39 |         nodes (list): list of nodes
 40 | 
 41 |     Returns:
 42 |         list: list of new nodes all with INT32 constants.
 43 |     """
 44 |     new_nodes = []
 45 |     for node in nodes:
 46 |         if (
 47 |             node.op_type == "Constant"
 48 |             and node.attribute[0].t.data_type == TensorProto.INT64
 49 |         ):
 50 |             data = nph.to_array(node.attribute[0].t).astype(np.int32)
 51 |             new_t = nph.from_array(data)
 52 |             new_node = h.make_node(
 53 |                 "Constant",
 54 |                 inputs=[],
 55 |                 outputs=node.output,
 56 |                 name=node.name,
 57 |                 value=new_t,
 58 |             )
 59 |             new_nodes += [new_node]
 60 |         else:
 61 |             new_nodes += [node]
 62 | 
 63 |     return new_nodes
 64 | 
 65 | 
 66 | def convert_model_to_int32(model_path: str, out_path: str):
 67 |     """
 68 |     convert_model_to_int32 Converts ONNX model with INT64 params to INT32 params.\n
 69 | 
 70 |     Args:\n
 71 |         model_path (str): path to original ONNX model.\n
 72 |         out_path (str): path to save converted model.
 73 |     """
 74 |     log.info("ONNX INT64 --> INT32 Converter")
 75 |     log.info(f"Loading Model: {model_path}")
 76 |     # * load model.
 77 |     model = onnx.load_model(model_path)
 78 |     ch.check_model(model)
 79 |     # * get model opset version.
 80 |     opset_version = model.opset_import[0].version
 81 |     graph = model.graph
 82 |     # * The initializer holds all non-constant weights.
 83 |     init = graph.initializer
 84 |     # * collect model params in a dictionary.
 85 |     params_dict = make_param_dictionary(init)
 86 |     log.info("Converting INT64 model params to INT32...")
 87 |     # * convert all INT64 aprams to INT32.
 88 |     converted_params = convert_params_to_int32(params_dict)
 89 |     log.info("Converting constant INT64 nodes to INT32...")
 90 |     new_nodes = convert_constant_nodes_to_int32(graph.node)
 91 | 
 92 |     graph_name = f"{graph.name}-int32"
 93 |     log.info("Creating new graph...")
 94 |     # * create a new graph with converted params and new nodes.
 95 |     graph_int32 = h.make_graph(
 96 |         new_nodes,
 97 |         graph_name,
 98 |         graph.input,
 99 |         graph.output,
100 |         initializer=converted_params,
101 |     )
102 |     log.info("Creating new int32 model...")
103 |     model_int32 = h.make_model(graph_int32, producer_name="onnx-typecast")
104 |     model_int32.opset_import[0].version = opset_version
105 |     ch.check_model(model_int32)
106 |     log.info(f"Saving converted model as: {out_path}")
107 |     onnx.save_model(model_int32, out_path)
108 |     log.info(f"Done Done London. 🎉")
109 |     return
110 | 
111 | 
112 | if __name__ == "__main__":
113 |     typer.run(convert_model_to_int32)
114 | 


--------------------------------------------------------------------------------