HarleyHK / onnx-tool

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

Geek Repo:Geek Repo

Github PK Tool:Github PK Tool

onnx-tool

A tool for ONNX model:

  • Rapid shape inference.
  • Profile model.
  • Compute Graph and Shape Engine.
  • OPs fusion.
  • Quantized models and sparse models are supported.

Supported Models:

  • NLP: BERT, T5, GPT
  • Diffusion: Stable Diffusion(TextEncoder, VAE, UNET)
  • CV: Resnet, MobileNet, YOLO, ...
  • Audio: LPCNet

Shape inference

how to use: data/Profile.md.
pytorch usage: data/PytorchUsage.md.
tensorflow usage: data/TensorflowUsage.md.
samples: benchmark/samples.py.


Profile Model

Float MultipleAdd Count(1 MAC=2 FLOPs), Memory Usage(in bytes), Parameters(elements number)

Sparse Pattern, Sparse Block Ratio, Sparse Element Ratio

how to use: data/Profile.md.
pytorch usage: data/PytorchUsage.md.
tensorflow usage: data/TensorflowUsage.md.
samples: benchmark/samples.py.


Compute Graph with Shape Engine

Remove shape calculation layers(created by ONNX export) to get a Compute Graph. Use Shape Engine to update tensor shapes at runtime.
Samples: benchmark/shape_regress.py. benchmark/samples.py.
Integrate Compute Graph and Shape Engine into a cpp inference engine: data/inference_engine.md


Inplace op fusion

MHA and Layernorm Fusion for Transformers

Resnet18 fusion

how to use: data/Subgraph.md.
BERT samples: benchmark/samples.py.


Extract subgraph from ONNX model

how to use: data/Subgraph.md.


Tensor operations

  • Export weight tensors to files
  • Simplify tensor and node names, convert name from a long string to a short string
  • Remove unused tensors, models like vgg19-7.onnx set its static weight tensors as its input tensors
  • Set custom input and output tensors' name and dimension, change model from fixed input to dynamic input
    how to use: data/Tensors.md.

How to install

pip install onnx-tool

OR

pip install --upgrade git+https://github.com/ThanatosShinji/onnx-tool.git

python>=3.6

If pip install onnx-tool failed by onnx's installation, you may try pip install onnx==1.8.1 (a lower version like this) first.
Then pip install onnx-tool again.


Known Issues

  • Loop op is not supported

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)
text_encoder 123.13 6,782
UNet2DCondition 859.52 888,870
VAE_encoder 34.16 566,371
VAE_decoder 49.49 1,271,959
SqueezeNet 1.0 1.23 351
VGG 19 143.66 19,643
AlexNet 60.96 665
GoogleNet 6.99 1,606
googlenet_age_adience 5.98 1,605
LResNet100E-IR 65.22 12,102
BERT-Squad 113.61 22,767
BiDAF 18.08 9.87
EfficientNet-Lite4 12.96 1,361
Emotion FERPlus 12.95 877
Mask R-CNN R-50-FPN-fp32 46.77 92,077
Model Params(M) MACs(M)
rvm_mobilenetv3_fp32.onnx 3.73 4,289
yolov4 64.33 3,319
ConvNeXt-L 229.79 34,872
edgenext_small 5.58 1,357
SSD 19.98 216,598
RealESRGAN_x4plus.pth 16.69 73,551
ShuffleNet-v2-fp32 2.29 146
GPT-2 137.02 1,103
T5-encoder 109.62 686
T5-decoder-with-lm-head 162.62 1,113
RoBERTa-BASE 124.64 688
Faster R-CNN R-50-FPN-fp32 44.10 46,018
FCN ResNet-50 35.29 37,056
MobileNet v2-1.0-fp32 3.3 300
ResNet50_fp32 25 3,868

About

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

License:MIT License


Languages

Language:Python 100.0%