/onnx-tool

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

Primary LanguagePythonMIT LicenseMIT

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