Package: onnx-tool
简体中文onnx-toolA tool for ONNX model:Parse and edit:Constant folding;OPs fusion.Model profiling: Rapid shape inference; MACs statisticsCompute Graph and Shape Engine.Model memory compression: activation compression and weight compression.Quantized models and sparse modelsare supported.Supported Models:NLP: BERT, T5, GPT, LLaMa, MPT(TransformerModel)Diffusion: Stable Diffusion(TextEncoder, VAE, UNET)CV:BEVFormer, MobileNet, YOLO, ...Audio: sovits, LPCNetBasic Parse and EditYou can load any onnx file by onnx_tool.Model:Change graph structure with onnx_tool.Graph;Change op attributes and IO tensors with onnx_tool.Node;Change tensor data or type with onnx_tool.Tensor.To apply your changes, just call save_model method of onnx_tool.Model or onnx_tool.Graph.Please referbenchmark/examples.py.Shape Inference & Profile ModelAll profiling data must be built on shape inference result.ONNX graph with tensor shapes:Regular model profiling table:Sparse profiling table:Introduction:data/Profile.md.pytorch usage:data/PytorchUsage.md.tensorflow
usage:data/TensorflowUsage.md.examples:benchmark/examples.py.Compute Graph with Shape EngineFrom a raw graph to a compute graph:Remove shape calculation layers(created by ONNX export) to get aCompute Graph. UseShape Engineto update tensor
shapes at runtime.Examples:benchmark/shape_regress.py.benchmark/examples.py.IntegrateCompute GraphandShape Engineinto a cpp inference
engine:data/inference_engine.mdMemory CompressionActivation CompressionActivation memory also called temporary memory is created by each OP's output. Only the last activation marked as the
model's output will be kept. So you don't have to prepare memory space for each activation tensor. They better reuse
an optimized memory size.For large language models and high-resolution CV models, the activation memory compression is a key to save memory.The compression method achieves 5% memory compression on most models.For example:modelNative Memory Size(MB)Compressed Memory Size(MB)Compression Ratio(%)StableDiffusion(VAE_encoder)14,2455403.7StableDiffusion(VAE_decoder)25,4171,1404.48StableDiffusion(Text_encoder)21552.5StableDiffusion(UNet)36,1352,2326.2GPT24026.9BERT2,170271.25code example:benchmark/compression.pyWeight CompressionA fp32 model with 7B parameters will take 28GB disk space and memory space. You can not even run the model if your device
doesn't have that much memory space. So weight compression is critical to run large language models. As a reference, 7B
model with int4 symmetric per block(32) quantization(llama.cpp's q4_0 quantization method) only has ~0.156x model size compared with fp32 model.Current support:[fp16][int8]x[symmetric/asymmetric]x[per tensor/per channel/per block][int4]x[symmetric/asymmetric]x[per tensor/per channel/per block]code examples:benchmark/examples.py.How to installpip install onnx-toolORpip install --upgrade git+https://github.com/ThanatosShinji/onnx-tool.gitpython>=3.6Ifpip install onnx-toolfailed by onnx's installation, you may trypip install onnx==1.8.1(a lower version like this) first.Thenpip install onnx-toolagain.Known IssuesLoop op is not supportedSequence type is not supportedResults ofONNX Model Zooand SOTA modelsSome models have dynamic input shapes. The MACs varies from input shapes. The input shapes used in these results are writen todata/public/config.py.
These onnx models with all tensors' shape can be downloaded:baidu drive(code: p91k)google driveModelParams(M)MACs(M)GPT-J 1 layer464173,398MPT 1 layer26179,894text_encoder123.136,782UNet2DCondition859.52888,870VAE_encoder34.16566,371VAE_decoder49.491,271,959SqueezeNet 1.01.23351AlexNet60.96665GoogleNet6.991,606googlenet_age5.981,605LResNet100E-IR65.2212,102BERT-Squad113.6122,767BiDAF18.089.87EfficientNet-Lite412.961,361Emotion12.95877Mask R-CNN46.7792,077ModelParams(M)MACs(M)LLaMa 1 layer618211,801BEVFormer Tiny33.7210,838rvm_mobilenetv33.734,289yolov464.333,319ConvNeXt-L229.7934,872edgenext_small5.581,357SSD19.98216,598RealESRGAN16.6973,551ShuffleNet2.29146GPT-2137.021,103T5-encoder109.62686T5-decoder162.621,113RoBERTa-BASE124.64688Faster R-CNN44.1046,018FCN ResNet-5035.2937,056ResNet50253,868