onnx-tool 0.1.8

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

onnx-tool

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

• Shape inference

---

• MACs counting for each node (FLOPs=2*MACs)

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_tool
  modelpath = 'resnet50-v1-12.onnx'
  onnx_tool.model_profile(modelpath, None, None) #pass file name
  

  import onnx
  import onnx_tool
  modelpath = 'resnet50-v1-12.onnx'
  model = onnx.load_model(modelpath)
  onnx_tool.model_shape_infer(model, None, saveshapesmodel='resnet50_shapes.onnx')  # pass ONNX.ModelProto
  

  cli usage (dynamic shapes is not supported)

  python -m onnx_tool -i 'resnet50-v1-12.onnx' -o 'resnet50_shapes.onnx'
  

• Dynamic input shapes and dynamic resize scales('downsample_ratio')

  import numpy
  import onnx_tool
  from onnx_tool import create_ndarray_f32 #or use numpy.ones(shape,numpy.float32) is ok
  modelpath = '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)}
  onnx_tool.model_profile(modelpath,inputs,None,saveshapesmodel='rvm_mobilenetv3_fp32_shapes.onnx')
  

• Define your custom op's node profiler.

  import numpy
  import onnx
  import onnx_tool
  from onnx_tool import NODEPROFILER_REGISTRY,NodeBase,create_ndarray_f32
  
  @NODEPROFILER_REGISTRY.register()
  class CropPlugin(NodeBase):
      def __init__(self,nodeproto:onnx.NodeProto):
          super().__init__(nodeproto)
          #parse your attributes here
  
      def infer_shape(self, intensors: list[numpy.ndarray]):
          #calculate output shapes here
          #this node crops intensors[0] to the shape of intensors[1], just return list of intensors[1]
          #no need to finish the true calculation, just return a ndarray of a right shape
          return [intensors[1]]
  
      def profile(self,intensors:list[numpy.ndarray],outtensors:list[numpy.ndarray]):
          macs=0
          params=0
          #accumulate macs and params here
          #this node has no calculation
          return macs,params
  
  onnx_tool.model_profile('./rrdb_new.onnx', {'input': create_ndarray_f32((1, 3, 335, 619))},
                          savenode='rrdb_new_nodemap.txt', saveshapesmodel='rrdb_new_shapes.onnx')
  

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. These onnx models with all tensors' shape can be downloaded: baidu drive(code: p91k) google drive

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