AITemplate is a deep learning compiler framework that transforms neural network models into optimized CUDA or HIP code for high-performance GPU inference. It features a frontend for model definition (similar to PyTorch), a compiler that applies graph optimizations (operator fusion, shape propagation), and backends that generate architecture-specific kernels. The framework supports model conversion from PyTorch via FX graph transformation and includes examples for vision (ResNet, ViT, Stable Diffusion), NLP (BERT), and object detection (Detectron2). Key innovations include a universal GEMM framework with flexible memory layouts, operator fusion patterns, and a tensor accessor abstraction for handling diverse memory layouts in generated code.
onnx2tf is a production-ready ONNX to TensorFlow/TFLite conversion framework that employs a sophisticated pass-based optimization pipeline. The system converts ONNX models through preprocessing, lowering to an intermediate representation (ModelIR), applying 50+ optimization passes (layout transformation, operator fusion, quantization), and exporting to target formats including TFLite, SavedModel, PyTorch nn.Module, and TorchScript. Its architecture features clean separation between operator implementations, IR transformation passes, and backend-specific exporters, with the flatbuffer_direct backend providing 100x+ faster conversion by bypassing TensorFlow graph construction. The framework includes comprehensive accuracy validation tools, dynamic shape handling, and model surgery capabilities for intermediate layer extraction.
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.