$ show ThanatosShinji/onnx-tool
This repository is a comprehensive toolkit for analyzing, optimizing, and transforming ONNX models. It provides a core library (onnx_tool) for loading models, performing graph analysis (shape/value inference), profiling computational complexity (MACs, parameters, memory), applying optimizations (operator fusion, constant folding), quantizing weights, and managing memory for efficient inference. The toolkit supports a wide range of model architectures including LLMs (BERT, GPT, LLaMa, MPT, Qwen), diffusion models (Stable Diffusion), computer vision (ResNet, EfficientNet), and multimodal models. Key features include a custom inference engine (GraphInfer) with memory pooling for zero-copy tensor views, a plugin-based node registry system for extending ONNX operator support, and utilities for model conversion and evaluation. The architecture centers around the onnx_tool package, which defines fundamental abstractions like Graph, Tensor, and Node. The Graph class handles model loading, shape inference, and transformation pipelines. The Node registry system allows dynamic registration of operation-specific handlers for shape/value inference and profiling. Optimization modules (fusion, quantization) work on the graph representation, while the inference engine (inference/infer.py) executes optimized graphs using a memory pool to minimize allocation overhead. Additional components include benchmarking scripts for performance analysis, model configuration management, and specialized support for LLM KV-cache handling. While the core library demonstrates solid engineering with reusable patterns like the node registry and graph rewrite framework, there are areas of technical drift due to loose dependency versioning and the broad scope requiring continuous maintenance to keep pace with evolving model architectures and ONNX specifications. The project appears maintained based on recent commit activity and community engagement (stars, forks, issue templates), though dependency management without version constraints poses a risk for long-term stability.
The LLM identified 2 files worth highlighting for reuse:
Implements a base Node class with shape inference, value inference, and profiling methods, along with a registry system (NODE_REGISTRY) for dynamically registering operation-specific handlers. Includes mixin classes (PWNode for pointwise operations, FusedBase for simple forwarding) to reduce code duplication. This plugin architecture enables extensible ONNX graph processing without modifying core code, making it highly reusable in any project requiring custom ONNX operator support.
Provides a framework for pattern-based graph rewrites to optimize ONNX models. Implements common fusion patterns (e.g., Conv-BN, residual connections, element-wise operations) using a declarative pattern-matching approach. The framework is decoupled from specific model architectures and allows easy extension with new patterns. This isolated graph transformation system demonstrates non-obvious skill in optimizing computation graphs and can be extracted for use in other ONNX optimization tools.