From f8537f7757e12e318b80d2cd488845b4c0b33e9c Mon Sep 17 00:00:00 2001
From: mansterteddy <mansterteddy@gmail.com>
Date: Thu, 14 Oct 2021 13:57:59 +0800
Subject: [PATCH] Add customized_op.

---
 onnxruntime/customized_op/export.py    | 27 ++++++++++++++++++++++++++
 onnxruntime/customized_op/lib_path.py  |  3 +++
 onnxruntime/customized_op/run_ort.py   | 17 ++++++++++++++++
 onnxruntime/customized_op/run_torch.py | 14 +++++++++++++
 4 files changed, 61 insertions(+)
 create mode 100644 onnxruntime/customized_op/export.py
 create mode 100644 onnxruntime/customized_op/lib_path.py
 create mode 100644 onnxruntime/customized_op/run_ort.py
 create mode 100644 onnxruntime/customized_op/run_torch.py

diff --git a/onnxruntime/customized_op/export.py b/onnxruntime/customized_op/export.py
new file mode 100644
index 0000000..db1d92f
--- /dev/null
+++ b/onnxruntime/customized_op/export.py
@@ -0,0 +1,27 @@
+import torch
+
+class CustomInverse(torch.nn.Module):
+    def forward(self, x):
+        return torch.inverse(x) + x 
+
+from torch.onnx import register_custom_op_symbolic
+
+def my_inverse(g, input):
+    return g.op("ai.onnx.contrib::Inverse", input)
+
+register_custom_op_symbolic('::inverse', my_inverse, 1)
+
+x0 = torch.Tensor([[-3.7806,  1.0857, -0.8645],
+        [-0.0398,  0.3996, -0.7268],
+        [ 0.3433,  0.6064,  0.0934]])
+t_model = CustomInverse()
+print(t_model(x0))
+
+torch.onnx.export(t_model, 
+                    (x0, ),
+                    "export.onnx", 
+                    opset_version=12, 
+                    verbose=True,
+                    input_names= ["input"],
+                    output_names= ["output"]
+                    )
\ No newline at end of file
diff --git a/onnxruntime/customized_op/lib_path.py b/onnxruntime/customized_op/lib_path.py
new file mode 100644
index 0000000..0a2ad03
--- /dev/null
+++ b/onnxruntime/customized_op/lib_path.py
@@ -0,0 +1,3 @@
+from onnxruntime_extensions import get_library_path as _lib_path
+
+print(_lib_path())
\ No newline at end of file
diff --git a/onnxruntime/customized_op/run_ort.py b/onnxruntime/customized_op/run_ort.py
new file mode 100644
index 0000000..f3c038b
--- /dev/null
+++ b/onnxruntime/customized_op/run_ort.py
@@ -0,0 +1,17 @@
+import onnx
+import torch
+import numpy as np
+import onnxruntime as ort
+from onnxruntime_extensions import onnx_op, PyOp, PyOrtFunction
+
+@onnx_op(op_type="Inverse")
+def inverse(x):
+        return np.linalg.inv(x)
+
+x = [[-3.7806,  1.0857, -0.8645], [-0.0398,  0.3996, -0.7268], [ 0.3433,  0.6064,  0.0934]]
+x = np.asarray(x, dtype=np.float32)
+
+onnx_model = onnx.load("export.onnx")
+onnx_fn = PyOrtFunction.from_model(onnx_model)
+y = onnx_fn(x)
+print(y)
\ No newline at end of file
diff --git a/onnxruntime/customized_op/run_torch.py b/onnxruntime/customized_op/run_torch.py
new file mode 100644
index 0000000..61d3ad2
--- /dev/null
+++ b/onnxruntime/customized_op/run_torch.py
@@ -0,0 +1,14 @@
+import torch
+
+class CustomInverse(torch.nn.Module):
+    def forward(self, x):
+        return torch.inverse(x) + x 
+
+x = torch.Tensor([[-3.7806,  1.0857, -0.8645],
+        [-0.0398,  0.3996, -0.7268],
+        [ 0.3433,  0.6064,  0.0934]])
+model = CustomInverse()
+y = model(x)
+print(y)
+print(x.dtype)
+#torch.onnx.export(model, (x), "export_torch.onnx")
\ No newline at end of file