PyTorch to ONNX Exporter

Experimental torch ONNX exporter. Compatible with torch>=2.1.

[!WARNING] This is an experimental project and is not designed for production use. Use torch.onnx.export for these purposes.

Installation

pip install --upgrade torch-onnx

Usage

import torch
import torch_onnx
from onnxscript import ir
import onnx

# Get an exported program with torch.export
exported = torch.export.export(...)
model = torch_onnx.exported_program_to_ir(exported)
proto = ir.to_proto(model)
onnx.save(proto, "model.onnx")

# Or patch the torch.onnx export API
# Set error_report=True to get a detailed error report if the export fails
torch_onnx.patch_torch(report=True, verify=True, profile=True)
torch.onnx.export(...)

# Use the analysis API to print an analysis report for unsupported ops
torch_onnx.analyze(exported)

Design

{dynamo/jit} -> {ExportedProgram} -> {torchlib} -> {ONNX IR} -> {ONNX}

Use ExportedProgram
- Rely on robustness of the torch.export implementation
- Reduce complexity in the exporter
- This does not solve dynamo limitations, but it avoids introducing additional breakage by running fx passes
Flat graph; Scope info as metadata, not functions
- Because existing tools are not good at handling them
Eager optimization where appropriate
- Because exsiting tools are not good at optimizing
Drop in replacement for torch.onnx.export
- Minimum migration effort
Build graph eagerly in the exporter
- Give the exporter full control over the graph being built

Why is this doable?

We need to verify torch.export coverage on Huggingface Optimum https://github.com/huggingface/optimum/tree/main/optimum/exporters/onnx; and they are not patching torch.onnx itself.
Patch torch.onnx.export such that packages do not need to change a single line to use dynamo
We have all operators implemented and portable

Metadata

Release files for torch-onnx 0.1.25

For a detailed explanation of source distributions (sdists) and built distributions (wheels), please see the package formats documentation.

Source distribution (sdist)

Source distribution for torch-onnx 0.1.25
File	Size	Uploaded
torch_onnx-0.1.25.tar.gz	72.8 kB	Sep 18, 2024	Details

Built distribution (wheel)

Table of built distributions (wheels) for torch-onnx 0.1.25
File	Interpreter	ABI	Platform	Reset
torch_onnx-0.1.25-py3-none-any.whl 81.6 kB Sep 18, 2024	Python 3	none	any	Details

Total release size: 154.5 kB

Release files / torch_onnx-0.1.25.tar.gz

Download URL	torch_onnx-0.1.25.tar.gz
Size	72.8 kB
Tags	Source
SHA-256 checksum How to use checksums	`cddcc6089cbfdd3e5ab65dc8ed93249e257215eb0f6c068ffa7377d6ef6b767e`
BLAKE2b-256 checksum How to use checksums	`87bc454c57e61e47875245230491313cd5813ce15eb3aa8633e8064008097924`
Upload date	Sep 18, 2024
Uploaded using Trusted Publishing? What is trusted publishing?	Yes
Uploaded via	`twine/4.0.1 CPython/3.11.10`

Release files / torch_onnx-0.1.25-py3-none-any.whl

Download URL	torch_onnx-0.1.25-py3-none-any.whl
Size	81.6 kB
Tags	Python 3
SHA-256 checksum How to use checksums	`a9ce6d5fe4d77c4e697cf97dbbf919a1b77834b299a9e9debd2a25d5415ce68d`
BLAKE2b-256 checksum How to use checksums	`9686a5bae457e4fceb61cf873789e59f265f83bac0380ffb607abba42ba472a5`
Upload date	Sep 18, 2024
Uploaded using Trusted Publishing? What is trusted publishing?	Yes
Uploaded via	`twine/4.0.1 CPython/3.11.10`