Skip to content

Repository files navigation

📄 简体中文 | ✨ New Project: AI-Enhancement-Filter (powered by onnx-tool)


Python 3.6+ PyPI Version License

onnx-tool

A comprehensive toolkit for analyzing, optimizing, and transforming ONNX models with advanced capabilities for LLMs, diffusion models, and computer vision architectures.

  • LLM Optimization: Build and profile large language models with KV cache analysis (example)
  • Graph Transformation:
    • Constant folding (docs)
    • Operator fusion (docs)
  • Advanced Profiling:
    • Rapid shape inference
    • MACs/parameter statistics with sparsity awareness
  • Compute Graph Engine: Runtime shape computation with minimal overhead (details)
  • Inference Engine: PyTorch-backed graph inference with memory pool (details)
    • 40+ registered op kernels (Conv, Add, Relu, Gemm, etc.)
    • Up to 11.3x faster than PyTorch in dynamic resolution scenarios
    • 54% memory reduction via two-pass compression algorithm
  • Memory Compression:
    • Activation memory optimization (up to 95% reduction)
    • Weight quantization (FP16, INT8/INT4 with per-tensor/channel/block schemes)
  • Quantization & Sparsity: Full support for quantized and sparse model analysis

🤖 Supported Model Architectures

Domain Models
NLP BERT, T5, GPT, LLaMa, MPT, Qwen3, Qwen3.5 (Dense & MoE), DeepSeek-V4 (Flash/Pro, MLA+MoE), MiniMax-M2.7 (MoE) (TransformerModel)
Diffusion Stable Diffusion (TextEncoder, VAE, UNet)
CV Detic, BEVFormer, SSD300_VGG16, ConvNeXt, Mask R-CNN, Silero VAD
Audio Sovits, LPCNet

🆕 Qwen3.5 Series: Full support for Qwen3.5 hybrid architecture including:

  • Gated DeltaNet (GDN) layers with linear attention
  • QKV Gating (Q projection with built-in gate, applied before O-projection)
  • Sparse Mixture-of-Experts (MoE) with routed + shared experts
  • Mixed layer types (linear_attention / full_attention) per config
  • Multimodal Vision Encoder (ViT + MLP Projector) with resolution-aware profiling

⚡ Build & Profile LLMs in Seconds

Profile 10 Hugging Face models in under one second. Export ONNX models with llama.cpp-like simplicity (code).

Model Statistics (1k token input)

model name(1k input) MACs(G) Parameters(G) KV Cache(G)
Phi-3-mini-4k 4083 3.82108 0.201327
Phi-3-small-8k-instruct 7912 7.80167 0.0671089
Phi-3-medium-4k-instruct 14665 13.9602 0.104858
Llama3-8B 8029 8.03026 0.0671089
Llama-3.1-70B-Japanese-Instruct-2407 72888 70.5537 0.167772
Qwen3.5-4B-Instruct 🆕 4807 4.651 0.067109
Qwen3.5-35B-A3B-Instruct 🆕 (MoE) 3574 34.705 0.041943
DeepSeek-V4-Flash 🆕 (MoE/MLA) 15681 283.811 0.045089
DeepSeek-V4-Pro 🆕 (MoE/MLA) 55701 1571.742 0.063963
MiniMax-M2.7 🆕 (MoE) 12554 230.315 0.130023

Multimodal: Vision Encoder vs LLM MACs 🆕

Vision encoder (ViT-24L + MLP Projector) vs LLM backbone. LLM input seq_len = vision patches + 1K text tokens. Both models share the same ViT architecture. See benchmark/vision_llm_compare.py.

Resolution Patches Total Tokens Vision(G) 4B LLM(G) 4B Vis/LLM(%) 35B LLM(G) 35B Vis/LLM(%)
224×224 256 1,280 84 6,080 1.4 4,557 1.8
448×448 1,024 2,048 389 10,070 3.9 7,719 5.0
672×672 2,304 3,328 1,061 17,292 6.1 13,702 7.7
896×896 4,096 5,120 2,350 28,599 8.2 23,576 9.9
1344×896 6,144 7,168 4,317 43,233 10.0 37,000 11.6
1344×1344 9,216 10,240 8,259 68,607 12.0 61,414 13.4

💡 Vision MACs scale with O(p²), LLM MACs with O(t²) where t = patches + text. At 448×448, vision is ~4-5% of LLM; at 1344×1344, ~12-13%. Vision is always the minor component.

MoE Activated Parameters vs Sequence Length

Activated parameters = sum of all nodes' static_params (weights actually accessed during forward pass). S=32: Qwen3.5-35B-A3B & MiniMax-M2.7 experts fully activated (gap ≈ embedding); DeepSeek-V4 not yet saturated (top-6 × 32 = 192 < num_experts).

model Total(G) S=1 S=2 S=4 S=8 S=16 S=32
Qwen3.5-35B-A3B (MoE) 34.3 2.44 3.44 5.46 9.48 17.54 33.64
MiniMax-M2.7 (MoE) 230.3 10.42 17.44 31.48 59.56 115.73 228.08
DeepSeek-V4-Flash (MoE) 283.8 12.75 19.24 32.23 58.20 110.14 214.03
DeepSeek-V4-Pro (MoE) 1571.7 47.60 71.77 120.13 216.84 410.27 797.12

Latency Estimation (4-bit weights, 16-bit KV cache)

Prefill Throughput (tokens/s, 1k input)

model Ultra-358H Arc-B70 RTX-4090 RTX-5090
Phi-3-mini-4k 6658.9 21279.8 19896.0 25567.0
Phi-3-small-8k-instruct 3599.5 11316.9 10429.8 13334.3
Phi-3-medium-4k-instruct 1964.8 6151.8 5649.1 7213.2
Llama3-8B 3499.9 11053.2 10226.7 13092.7
Llama-3.1-70B-Japanese-Instruct-2407 401.2 1249.3 1141.8 1455.5
Qwen3.5-4B-Instruct 🆕 5705.0 18162.0 16923.1 21720.3
Qwen3.5-35B-A3B-Instruct 🆕 (MoE) 4660.9 18458.3 21970.0 29367.3
MiniMax-M2.7 🆕 (MoE) 970.5 4285.1 6090.4 9083.3

Decode Throughput (tokens/s)

model Ultra-358H Arc-B70 RTX-4090 RTX-5090
Phi-3-mini-4k 56.4 267.9 444.2 789.7
Phi-3-small-8k-instruct 33.4 158.5 262.8 467.2
Phi-3-medium-4k-instruct 18.0 85.4 141.6 251.8
Llama3-8B 32.9 156.1 258.9 460.2
Llama-3.1-70B-Japanese-Instruct-2407 3.6 17.3 28.7 51.0
Qwen3.5-4B-Instruct 🆕 56.5 268.4 444.9 791.0
Qwen3.5-35B-A3B-Instruct 🆕 (MoE) 79.0 375.0 621.8 1105.4
MiniMax-M2.7 🆕 (MoE) 23.4 111.0 184.1 327.3

💡 Latencies computed from hardware specs – no actual inference required. Uses BF16/FP16 compute with FP32 accumulate as the standard.


🔧 Basic Parsing & Editing

Intuitive API for model manipulation:

from onnx_tool import Model

model = Model('model.onnx')          # Load any ONNX file
graph = model.graph                  # Access computation graph
node = graph.nodemap['Conv_0']       # Modify operator attributes
tensor = graph.tensormap['weight']   # Edit tensor data/types
model.save_model('modified.onnx')    # Persist changes

See comprehensive examples in benchmark/examples.py.


📊 Shape Inference & Profiling

All profiling relies on precise shape inference:

Shape inference visualization

Profiling Capabilities

  • Standard profiling: MACs, parameters, memory footprint
  • Sparse-aware profiling: Quantify sparsity impact on compute

MACs profiling table Sparse model profiling

📚 Learn more:


⚙️ Compute Graph & Shape Engine

Transform exported ONNX graphs into efficient Compute Graphs by removing shape-calculation overhead:

Compute graph transformation

  • Compute Graph: Minimal graph containing only compute operations
  • Shape Engine: Runtime shape resolver for dynamic models

Use Cases:

  • Integration with custom inference engines (guide)
  • Shape regression testing (example)

💾 Memory Compression

Activation Memory Compression

Reuses temporary buffers to minimize peak memory usage – critical for LLMs and high-res CV models.

model Native Memory Size(MB) Compressed Memory Size(MB) Compression Ratio(%)
StableDiffusion(VAE_encoder) 14,245 540 3.7
StableDiffusion(VAE_decoder) 25,417 1,140 4.48
StableDiffusion(Text_encoder) 215 5 2.5
StableDiffusion(UNet) 36,135 2,232 6.2
GPT2 40 2 6.9
BERT 2,170 27 1.25

✅ Typical models achieve >90% activation memory reduction
📌 Implementation: benchmark/compression.py

The compress_memory() algorithm has been patched with two improvements (see onnx_tool/graph.py):

  1. Size-sorted allocation: New tensors within each node are allocated in descending size order, reducing fragmentation
  2. Tail compression: Unused space at the end of the memory pool is trimmed
  3. List reference fix: Each tensor's [offset, size] is stored as an independent copy, preventing accidental cross-tensor aliasing

Benchmark results across models:

model Native(MB) Compressed(MB) Ratio(%)
VAE encoder 11,313.6 512.0 4.53
VAE decoder 19,816.2 896.1 4.52
Text encoder 172.5 3.7 2.12
GPT2 381.1 16.1 4.23
ResNet50 279.3 10.7 3.84

✅ Optimized algorithm achieves up to 54% additional pool reduction (ResNet50: 21.4→10.7 MB vs original)

Inference Engine

The inference/ module provides a complete PyTorch-backed inference engine built on the compressed memory pool:

  • MemoryPool: Zero-copy tensor views into a pre-allocated contiguous buffer
  • Kernel Registry: 40+ registered op kernels (Conv, Add, Relu, Gemm, MatMul, etc.)
  • GraphInfer: Graph-level inference with shape engine integration

Performance highlights (ResNet18 on Intel XPU):

Metric PyTorch GraphInfer Improvement
Single inference (1080p) 0.0151s 0.0167s 0.91x (on par)
Sequential 7 resolutions 1820ms 162ms 11.3x faster
Peak XPU memory (4K) 2338 MB 1926 MB 17.6% less
Memory pool (4K) 1012 MB Fixed size

📌 See inference/README.md for full benchmark details

Weight Compression

Essential for deploying large models on memory-constrained devices:

Quantization Scheme Size vs FP32 Example (7B model)
FP32 (baseline) 1.00× 28 GB
FP16 0.50× 14 GB
INT8 (per-channel) 0.25× 7 GB
INT4 (block=32, symmetric) – llama.cpp 0.156× 4.4 GB

Supported schemes:

  • ✅ FP16
  • ✅ INT8: symmetric/asymmetric × per-tensor/channel/block
  • ✅ INT4: symmetric/asymmetric × per-tensor/channel/block

📌 See benchmark/examples.py for implementation examples.


🚀 Installation

# PyPI (recommended)
pip install onnx-tool

# Latest development version
pip install --upgrade git+https://github.com/ThanatosShinji/onnx-tool.git

Requirements: Python ≥ 3.6

⚠️ Troubleshooting: If ONNX installation fails, try:

pip install onnx==1.8.1 && pip install onnx-tool

Known Issues

  • Loop op is not supported
  • Sequence type is not supported

📈 Model Zoo Results

Comprehensive profiling of ONNX Model Zoo and SOTA models. Input shapes defined in data/public/config.py.

📥 Download pre-profiled models (with full tensor shapes):

Model Params(M) MACs(M)
GPT-J 1 layer 464 173,398
MPT 1 layer 261 79,894
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
AlexNet 60.96 665
GoogleNet 6.99 1,606
googlenet_age 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 12.95 877
Mask R-CNN 46.77 92,077
Model Params(M) MACs(M)
LLaMa 1 layer 618 211,801
BEVFormer Tiny 33.7 210,838
rvm_mobilenetv3 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 16.69 73,551
ShuffleNet 2.29 146
GPT-2 137.02 1,103
T5-encoder 109.62 686
T5-decoder 162.62 1,113
RoBERTa-BASE 124.64 688
Faster R-CNN 44.10 46,018
FCN ResNet-50 35.29 37,056
ResNet50 25 3,868

🤝 Contributing

Contributions are welcome! Please open an issue or PR for:

  • Bug reports
  • Feature requests
  • Documentation improvements
  • New model support

About

A tool for parsing, editing, optimizing, and profiling ONNX models.

Topics

Resources

Contributing

Stars

492 stars

Watchers

8 watching

Forks

Releases

Used by

Contributors

Languages