A tool for ONNX model's shape inference and MACs counting.
Project description
onnx-tool
A tool for ONNX model's shape inference and MACs counting.
- Shape inference
- MACs counting for each node
How to install
pip install onnx-tool
OR
pip install --upgrade git+https://github.com/ThanatosShinji/onnx-tool.git
How to use
-
Basic usage
import onnx from onnx_tool.node_profilers import graph_profile,print_node_map model = onnx.load('resnet50.onnx') macs, params, node_map = graph_profile(model.graph, None) #shape inference included print_node_map(node_map) onnx.save_model(model,'resnet50_shapes.onnx') #save model with inferred shapes
-
Dynamic input shapes and dynamic resize scales('downsample_ratio')
import onnx from onnx_tool.node_profilers import graph_profile,print_node_map,create_ndarray_f32 model = onnx.load('rvm_mobilenetv3_fp32.onnx') inputs= {'src': create_ndarray_f32((1, 3, 1080, 1920)), 'r1i': create_ndarray_f32((1, 16, 135, 240)), 'r2i':create_ndarray_f32((1,20,68,120)),'r3i':create_ndarray_f32((1,40,34,60)), 'r4i':create_ndarray_f32((1,64,17,30)),'downsample_ratio':numpy.array((0.25,),dtype=numpy.float32)} macs, params, node_map = graph_profile(model.graph, inputs) #shape inference included print_node_map(node_map,'rvm_nodemap.txt') #save node map to file onnx.save_model(model,'rvm_mobilenetv3_fp32_shapes.onnx')
-
Define your custom op's node profiler.
from onnx_tool.node_profilers import graph_profile,NODEPROFILER_REGISTRY @NODEPROFILER_REGISTRY.register() class YourOp(): def __init__(self,nodeproto): #parse your attributes here def infer_shape(self,intensors:List[numpy.ndarray]): #calculate output shapes here #return a list of ndarray return outtensors def profile(self,intensors:List[numpy.ndarray],outtensors:List[numpy.ndarray]): #do macs and params accumulations here return macs,params macs, params, node_map = graph_profile(yourmodel.graph, None)
Known Issues
- Loop op is not supported
- Shared weight tensor will be counted more than once
Results of ONNX Model Zoo and SOTA models
Some models have dynamic input shapes. The MACs varies from input shapes. The input shapes used in these results are writen to data/public/config.py.
|
|
Project details
Release history Release notifications | RSS feed
Download files
Download the file for your platform. If you're not sure which to choose, learn more about installing packages.
Source Distribution
onnx-tool-0.1.3.tar.gz
(15.7 kB
view hashes)
Built Distribution
onnx_tool-0.1.3-py3-none-any.whl
(14.4 kB
view hashes)
Close
Hashes for onnx_tool-0.1.3-py3-none-any.whl
Algorithm | Hash digest | |
---|---|---|
SHA256 | 3456aa4caa625b3c2597a638bd0d8b5732fae9828507edbf11f4a32e721006f9 |
|
MD5 | 5b8137f4985e4555a77aaa038ca946f3 |
|
BLAKE2b-256 | 15216e3e57aee67400284334908c6441dc5d4f45a21f32b16ff0a2f3507f598e |