https://github.com/justinchuby/torch-onnx
A standalone version of the next PyTorch ONNX exporter
https://github.com/justinchuby/torch-onnx
deep-learning model-export onnx pytorch
Last synced: about 2 months ago
JSON representation
A standalone version of the next PyTorch ONNX exporter
- Host: GitHub
- URL: https://github.com/justinchuby/torch-onnx
- Owner: justinchuby
- License: mit
- Created: 2024-05-10T02:30:36.000Z (about 1 year ago)
- Default Branch: main
- Last Pushed: 2024-08-22T20:05:06.000Z (10 months ago)
- Last Synced: 2024-08-23T09:09:02.137Z (10 months ago)
- Topics: deep-learning, model-export, onnx, pytorch
- Language: Python
- Homepage:
- Size: 563 KB
- Stars: 2
- Watchers: 2
- Forks: 1
- Open Issues: 24
-
Metadata Files:
- Readme: README.md
- License: LICENSE
Awesome Lists containing this project
README
# PyTorch to ONNX Exporter
[](https://badge.fury.io/py/torch-onnx)
Experimental torch ONNX exporter. Compatible with torch>=2.1,<2.6.
> [!WARNING]
> This is an experimental project and is not designed for production use.
> Use `torch.onnx.export(..., dynamo=True)` for these purposes.
> The main logic from this project has been merged into PyTorch.
## Installation
```bash
pip install --upgrade torch-onnx
```
## Usage
```python
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