TorchONNX

Compile ONNX models into standalone PyTorch code (.py module + .pth state dict) with zero ONNX dependencies at inference time. Generated code supports torch.vmap for batched verification.

Installation

git clone https://github.com/ZhongkuiMa/torchonnx.git
cd torchonnx
pip install -e ".[dev]"

Requirements: Python 3.11+, PyTorch, onnx, onnxruntime, numpy

Quick Start

from torchonnx import TorchONNX

converter = TorchONNX(verbose=True)
converter.convert("model.onnx", target_py_path="model.py", target_pth_path="model.pth")

Load and run the generated model:

import torch
from model import Model

model = Model()
model.load_state_dict(torch.load("model.pth"))
output = model(torch.randn(1, 3, 32, 32))

Usage Guide

vmap-compatible output

By default (vmap_mode=True), TorchONNX replaces in-place ops with functional equivalents and .item() calls with tensor-based indexing, making the generated code compatible with torch.vmap:

model = Model()
model.load_state_dict(torch.load("model.pth"))
batch_outputs = torch.vmap(model)(torch.randn(100, 3, 32, 32))

Disable with vmap_mode=False if vmap compatibility is not needed.

Preprocessing

Use TorchONNX.preprocess() to load, validate, and normalize an ONNX model without generating code:

model_proto = TorchONNX.preprocess("model.onnx", target_opset=20)

Benchmark presets

BENCHMARKS_WITHOUT_BATCH_DIM is a set of benchmark names that lack a batch dimension. Use if_has_batch_dim(benchmark_name) to check whether a model from a given benchmark includes one:

from torchonnx import BENCHMARKS_WITHOUT_BATCH_DIM, if_has_batch_dim

has_batch = if_has_batch_dim("acasxu")

API reference

Symbol	Description
TorchONNX(verbose=False, use_shapeonnx=False)	Create converter instance
.convert(onnx_path, benchmark_name, target_py_path, target_pth_path, vmap_mode=True)	Compile ONNX to PyTorch (.py + .pth)
.preprocess(onnx_path, target_opset, infer_shapes, clear_docstrings)	Load, validate, and normalize ONNX model (static method)
BENCHMARKS_WITHOUT_BATCH_DIM	Set of benchmark names without a batch dimension
if_has_batch_dim(benchmark_name)	Return whether the benchmark model includes a batch dimension

Compiler Pipeline

TorchONNX is a 5-stage compiler:

Stage	Directory	Purpose
1. Normalize	normalize/	Load, validate, convert opset, infer shapes
2. Build	build/	Extract structural IR (ModelIR, NodeIR) from ONNX graph
3. Analyze	analyze/	Classify tensors, map ONNX ops to PyTorch types
4. Generate	generate/	Emit __init__, forward(), state dict, imports
5. Simplify	simplify/	Remove unused buffers, default args, format code

Supported Operations

Category	Operations
Layers	Conv1d/2d, ConvTranspose1d/2d, MaxPool2d, AvgPool2d, AdaptiveAvgPool2d, BatchNorm2d, Linear, Dropout, Upsample
Activations	ReLU, LeakyReLU, Sigmoid, Tanh, Softmax, ELU, GELU
Functions	F.pad, F.interpolate, torch.cat, torch.gather, torch.mean/sum/min/max/argmax, torch.clamp, torch.where, torch.full, torch.arange
Operators	add, sub, mul, div, matmul, pow, neg, equal
Tensor ops	flatten, reshape, permute, squeeze, unsqueeze, expand, slice, split, cast, sign, atan, cos, sin, floor, round
Quantization	QuantizeLinear, DequantizeLinear (scalar and per-axis)
ORT fusions	BiasGelu, LayerNormalization, SimplifiedLayerNormalization, SkipLayerNormalization

Transformer architectures are decomposed into these basic operations.

Generated modules preserve ONNX forward values, shapes, and dtypes and depend only on PyTorch at runtime. Operators without an exact pure-PyTorch conversion, including fused QLinear kernels, are rejected instead of delegated to ONNX Runtime or replaced with search-specific semantics.

Project Structure

torchonnx/
├── src/torchonnx/
│   ├── _torchonnx.py     # TorchONNX class (main API)
│   ├── presets.py        # Benchmark-specific settings
│   ├── normalize/        # Stage 1: ONNX normalization
│   ├── build/            # Stage 2: Structural IR
│   ├── analyze/          # Stage 3: Semantic IR + type mapping
│   ├── generate/         # Stage 4: Code generation + handlers
│   └── simplify/         # Stage 5: Code optimization
└── tests/
    ├── test_units/       # Unit tests
    └── test_benchmarks/  # VNNCOMP 2024 benchmark validation

Tests

pytest tests/ -v                              # all tests
pytest tests/test_units/test_torchonnx/ -v   # unit tests
pytest tests/test_benchmarks/ -v             # benchmark tests

Benchmark tests require vnncomp_benchmarks cloned as a sibling directory.

Contributing

See CONTRIBUTING.md.

License

MIT License — see LICENSE.