onnx-tool 0.1.2

A tool for ONNX model's shape inference and MACs counting.

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.

Model Params(M) MACs(M) MobileNet v2-1.0-fp32 3.3 300 ResNet50_fp32 25 3868 SqueezeNet 1.0 1.23 351 VGG 19 143.66 19643 AlexNet 60.96 665 GoogleNet 6.99 1606 googlenet_age_adience 5.98 1605 LResNet100E-IR 65.22 12102 BERT-Squad 113.61 22767 BiDAF 18.08 9.87 EfficientNet-Lite4 12.96 1361 Emotion FERPlus 12.95 877 Mask R-CNN R-50-FPN-fp32 46.77 92077

Model Params(M) MACs(M) rvm_mobilenetv3_fp32.onnx 3.73 4289 yolov4 64.33 33019 ConvNeXt-L 229.79 34872 edgenext_small 5.58 1357 SSD 19.98 216598 RealESRGAN_x4plus.pth 16.69 73551 ShuffleNet-v2-fp32 2.29 146 GPT-2 137.02 1103 T5-encoder 109.62 686 T5-decoder-with-lm-head 162.62 1113 RoBERTa-BASE 124.64 688 Faster R-CNN R-50-FPN-fp32 44.10 46018 FCN ResNet-50 35.29 37056