mirror of
https://github.com/saymrwulf/onnxruntime.git
synced 2026-07-14 18:12:05 +00:00
108 lines
3.8 KiB
Markdown
108 lines
3.8 KiB
Markdown
---
|
|
title: Export PyTorch model
|
|
nav_exclude: true
|
|
---
|
|
|
|
## Export PyTorch model with custom ONNX operators
|
|
{: .no_toc }
|
|
|
|
This document explains the process of exporting PyTorch models with custom ONNX Runtime ops. The aim is to export a PyTorch model with operators that are not supported in ONNX, and extend ONNX Runtime to support these custom ops.
|
|
|
|
## Contents
|
|
{: .no_toc }
|
|
|
|
* TOC placeholder
|
|
{:toc}
|
|
|
|
### Export Built-In Contrib Ops
|
|
|
|
"Contrib ops" refers to the set of custom ops that are built in to most ORT packages.
|
|
Symbolic functions for all contrib ops should be defined in [pytorch_export_contrib_ops.py](https://github.com/microsoft/onnxruntime/blob/master/onnxruntime/python/tools/pytorch_export_contrib_ops.py).
|
|
|
|
To export using those contrib ops, call `pytorch_export_contrib_ops.register()` before calling `torch.onnx.export()`. For example:
|
|
|
|
```python
|
|
from onnxruntime.tools import pytorch_export_contrib_ops
|
|
import torch
|
|
|
|
pytorch_export_contrib_ops.register()
|
|
torch.onnx.export(...)
|
|
```
|
|
|
|
### Export a Custom Op
|
|
|
|
To export a custom op that's not a contrib op, or that's not already included in `pytorch_export_contrib_ops`, one will need to
|
|
write and register a custom op symbolic function.
|
|
|
|
We take the Inverse operator as an example:
|
|
|
|
```python
|
|
from torch.onnx import register_custom_op_symbolic
|
|
|
|
def my_inverse(g, self):
|
|
return g.op("com.microsoft::Inverse", self)
|
|
|
|
# register_custom_op_symbolic('<namespace>::inverse', my_inverse, <opset_version>)
|
|
register_custom_op_symbolic('::inverse', my_inverse, 1)
|
|
```
|
|
|
|
`<namespace>` is a part of the torch operator name. For standard torch operators, namespace can be omitted.
|
|
|
|
`com.microsoft` should be used as the custom opset domain for ONNX Runtime ops. You can choose the custom opset version during op registration.
|
|
|
|
For more on writing a symbolic function, see the [torch.onnx documentation](https://pytorch.org/docs/master/onnx.html#adding-support-for-operators).
|
|
|
|
### Extend ONNX Runtime with Custom Ops
|
|
|
|
The next step is to add an op schema and kernel implementation in ONNX Runtime.
|
|
See [custom operators](../reference/operators/add-custom-op.md) for details.
|
|
|
|
### Test End-to-End: Export and Run
|
|
|
|
Once the custom op is registered in the exporter and implemented in ONNX Runtime, you should be able to export it and run it with ONNX Runtime.
|
|
|
|
Below you can find a sample script for exporting and running the inverse operator as part of a model.
|
|
|
|
The exported model includes a combination of ONNX standard ops and the custom ops.
|
|
|
|
This test also compares the output of PyTorch model with ONNX Runtime outputs to test both the operator export and implementation.
|
|
|
|
```python
|
|
import io
|
|
import numpy
|
|
import onnxruntime
|
|
import torch
|
|
|
|
|
|
class CustomInverse(torch.nn.Module):
|
|
def forward(self, x):
|
|
return torch.inverse(x) + x
|
|
|
|
x = torch.randn(3, 3)
|
|
|
|
# Export model to ONNX
|
|
f = io.BytesIO()
|
|
torch.onnx.export(CustomInverse(), (x,), f)
|
|
|
|
model = CustomInverse()
|
|
pt_outputs = model(x)
|
|
|
|
# Run the exported model with ONNX Runtime
|
|
ort_sess = onnxruntime.InferenceSession(f.getvalue())
|
|
ort_inputs = dict((ort_sess.get_inputs()[i].name, input.cpu().numpy()) for i, input in enumerate((x,)))
|
|
ort_outputs = ort_sess.run(None, ort_inputs)
|
|
|
|
# Validate PyTorch and ONNX Runtime results
|
|
numpy.testing.assert_allclose(pt_outputs.cpu().numpy(), ort_outputs[0], rtol=1e-03, atol=1e-05)
|
|
```
|
|
|
|
By default, the opset version will be set to `1` for custom opsets. If you'd like to export your
|
|
custom op to a higher opset version, you can specify the custom opset domain and version using
|
|
the `custom_opsets argument` when calling the export API. Note that this is different than the opset
|
|
version associated with default `ONNX` domain.
|
|
|
|
```python
|
|
torch.onnx.export(CustomInverse(), (x,), f, custom_opsets={"com.microsoft": 5})
|
|
```
|
|
|
|
Note that you can export a custom op to any version >= the opset version used at registration.
|