torch-to-onnx

Convert PyTorch models to ONNX using the modern torch.export API and the onnx_ir library.

Features

• ✅ Based on torch.export (FX graph) for accurate tracing.
• ✅ Exports to ONNX with full control over opset and IR versions.
• ✅ Lightweight and extensible – add your own ATen mappings.
• ✅ Produces standard ONNX models that can be loaded by any ONNX runtime.

Installation

pip install torch-to-onnx

If you haven't installed onnx-ir separately, you may need to:

pip install onnx-ir   # or follow your own onnx_ir installation

Usage

Basic Example

import torch
import torchvision.models as models
from torch_to_onnx import convert_exported_program_to_onnx
import onnx

# 1. Create a model
model = models.resnet18(pretrained=True).eval()
dummy_input = torch.randn(1, 3, 224, 224)

# 2. Export using torch.export
ep = torch.export.export(model, (dummy_input,), strict=True)

# 3. Convert to ONNX (returns an ir.Model)
ir_model = convert_exported_program_to_onnx(
    ep,
    opset_version=18,
    ir_version=13,
)

# 4. Serialize and save
ir.save(ir_model, "resnet18.onnx")
print("ONNX model saved!")

Customizing the Conversion

You can extend the operator mapping by editing ATEN_TO_ONNX_OP in the source code, or you can fork the repository and modify the converter.

API Reference

convert_exported_program_to_onnx(ep, opset_version=18, ir_version=13, remove_unused=True)

• ep: torch.export.ExportedProgram – the exported program.
• opset_version: int – ONNX opset version (default 18).
• ir_version: int – ONNX IR version (default 13).
• remove_unused: bool – currently a placeholder, does nothing.
• Returns: onnx_ir.Model – an in‑memory ONNX model representation.

Supported Operators

The converter currently supports common ATen operators such as:

• conv2d, batch_norm, relu, max_pool2d, avg_pool2d, add, sub, mul, div, matmul, linear, reshape, transpose, cat, softmax, sigmoid, tanh, dropout, slice, unsqueeze, squeeze, shape, etc.

If you encounter an unsupported operator, the converter will print a warning and skip it. You can easily add new mappings in the ATEN_TO_ONNX_OP dictionary.

Dependencies

• PyTorch >= 1.12.0
• ONNX >= 1.12.0
• onnx-ir (or your own fork of onnx_ir)

License

MIT

Author

zhengankun (ankun.zheng@qq.com)