Metadata-Version: 2.4
Name: torch-to-onnx
Version: 0.0.2
Summary: A tool for converting PyTorch models to ONNX
Author-email: zhengankun <ankun.zheng@qq.com>
License: MIT
Requires-Python: >=3.7
Description-Content-Type: text/markdown
License-File: LICENSE
Requires-Dist: torch>=1.9.0
Requires-Dist: onnx>=1.12.0
Requires-Dist: onnx-ir
Dynamic: license-file

# 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

```bash
pip install torch-to-onnx
```

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

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

## Usage

### Basic Example

```python
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)
