GitVita
New repositoryImport repositoryNew issue
Your repositories Your dashboard Sign in

microsoft/onnxruntime-extensions

Public

mirrored fromhttps://github.com/microsoft/onnxruntime-extensionsAvailable

Watch0Fork0Star0
CodeCommitsIssuesPull requestsActionsInsightsSecurity
4a0f8929494fa301baa6c59f617cce7872a7c4c8

Branches

Tags

  • No tags available.
0Branches0Tags
Go to file
Add file
Code

Clone

HTTPS

Download ZIP

onnxruntime-extensions/test

test/test_pyops.py

187lines · modecode

RawDownload
Latest commit unavailable.
1import os
2import unittest
3import numpy as np
4from numpy.testing import assert_almost_equal
5from onnx import helper, onnx_pb as onnx_proto
6import onnxruntime as _ort
7from onnxruntime_customops import (
8 onnx_op, PyCustomOpDef,
9 get_library_path as _get_library_path)
10
11
12def _create_test_model_test():
13 nodes = []
14 nodes.append(helper.make_node(
15 'CustomOpOne', ['input_1', 'input_2'], ['output_1'],
16 domain='ai.onnx.contrib'))
17 nodes.append(helper.make_node(
18 'CustomOpTwo', ['output_1'], ['output'],
19 domain='ai.onnx.contrib'))
20
21 input0 = helper.make_tensor_value_info(
22 'input_1', onnx_proto.TensorProto.FLOAT, [3, 5])
23 input1 = helper.make_tensor_value_info(
24 'input_2', onnx_proto.TensorProto.FLOAT, [3, 5])
25 output0 = helper.make_tensor_value_info(
26 'output', onnx_proto.TensorProto.INT32, [3, 5])
27
28 graph = helper.make_graph(nodes, 'test0', [input0, input1], [output0])
29 model = helper.make_model(
30 graph, opset_imports=[helper.make_operatorsetid('ai.onnx.contrib', 1)])
31 return model
32
33
34def _create_test_model():
35 nodes = []
36 nodes[0:] = [helper.make_node('Identity', ['input_1'], ['identity1'])]
37 nodes[1:] = [helper.make_node('PyReverseMatrix',
38 ['identity1'], ['reversed'],
39 domain='ai.onnx.contrib')]
40
41 input0 = helper.make_tensor_value_info(
42 'input_1', onnx_proto.TensorProto.FLOAT, [None, 2])
43 output0 = helper.make_tensor_value_info(
44 'reversed', onnx_proto.TensorProto.FLOAT, [None, 2])
45
46 graph = helper.make_graph(nodes, 'test0', [input0], [output0])
47 model = helper.make_model(
48 graph, opset_imports=[helper.make_operatorsetid('ai.onnx.contrib', 1)])
49 return model
50
51
52def _create_test_model_double(prefix, domain='ai.onnx.contrib'):
53 nodes = []
54 nodes[0:] = [helper.make_node('Identity', ['input_1'], ['identity1'])]
55 nodes[1:] = [helper.make_node('%sAddEpsilon' % prefix,
56 ['identity1'], ['customout'],
57 domain=domain)]
58
59 input0 = helper.make_tensor_value_info(
60 'input_1', onnx_proto.TensorProto.DOUBLE, [None, None])
61 output0 = helper.make_tensor_value_info(
62 'customout', onnx_proto.TensorProto.DOUBLE, [None, None])
63
64 graph = helper.make_graph(nodes, 'test0', [input0], [output0])
65 model = helper.make_model(
66 graph, opset_imports=[helper.make_operatorsetid(domain, 1)])
67 return model
68
69
70def _create_test_model_2outputs(prefix, domain='ai.onnx.contrib'):
71 nodes = [
72 helper.make_node('Identity', ['x'], ['identity1']),
73 helper.make_node(
74 '%sNegPos' % prefix, ['identity1'], ['neg', 'pos'],
75 domain=domain)
76 ]
77
78 input0 = helper.make_tensor_value_info(
79 'x', onnx_proto.TensorProto.FLOAT, [])
80 output1 = helper.make_tensor_value_info(
81 'neg', onnx_proto.TensorProto.FLOAT, [])
82 output2 = helper.make_tensor_value_info(
83 'pos', onnx_proto.TensorProto.FLOAT, [])
84
85 graph = helper.make_graph(nodes, 'test0', [input0], [output1, output2])
86 model = helper.make_model(
87 graph, opset_imports=[helper.make_operatorsetid(domain, 1)])
88 return model
89
90
91class TestPythonOp(unittest.TestCase):
92
93 @classmethod
94 def setUpClass(cls):
95
96 @onnx_op(op_type="PyReverseMatrix")
97 def reverse_matrix(x):
98 # The user custom op implementation here.
99 return np.flip(x, axis=0).astype(np.float32)
100
101 @onnx_op(op_type="PyAddEpsilon",
102 inputs=[PyCustomOpDef.dt_double],
103 outputs=[PyCustomOpDef.dt_double])
104 def add_epsilon(x):
105 # The user custom op implementation here.
106 return x + 1e-3
107
108 @onnx_op(op_type="PyNegPos",
109 inputs=[PyCustomOpDef.dt_float],
110 outputs=[PyCustomOpDef.dt_float, PyCustomOpDef.dt_float])
111 def negpos(x):
112 neg = x.copy()
113 pos = x.copy()
114 neg[x > 0] = 0
115 pos[x < 0] = 0
116 return neg, pos
117
118 def test_python_operator(self):
119 so = _ort.SessionOptions()
120 so.register_custom_ops_library(_get_library_path())
121 onnx_model = _create_test_model()
122 self.assertIn('op_type: "PyReverseMatrix"', str(onnx_model))
123 sess = _ort.InferenceSession(onnx_model.SerializeToString(), so)
124 input_1 = np.array(
125 [1, 2, 3, 4, 5, 6]).astype(np.float32).reshape([3, 2])
126 txout = sess.run(None, {'input_1': input_1})
127 assert_almost_equal(txout[0], np.array([[5., 6.], [3., 4.], [1., 2.]]))
128
129 def test_add_epsilon_python(self):
130 so = _ort.SessionOptions()
131 so.register_custom_ops_library(_get_library_path())
132 onnx_model = _create_test_model_double('Py')
133 self.assertIn('op_type: "PyAddEpsilon"', str(onnx_model))
134 sess = _ort.InferenceSession(onnx_model.SerializeToString(), so)
135 input_1 = np.array([[0., 1., 1.5], [7., 8., -5.5]])
136 txout = sess.run(None, {'input_1': input_1})
137 diff = txout[0] - input_1 - 1e-3
138 assert_almost_equal(diff, np.zeros(diff.shape))
139
140 def test_python_negpos(self):
141 so = _ort.SessionOptions()
142 so.register_custom_ops_library(_get_library_path())
143 onnx_model = _create_test_model_2outputs('Py')
144 self.assertIn('op_type: "PyNegPos"', str(onnx_model))
145 sess = _ort.InferenceSession(onnx_model.SerializeToString(), so)
146 x = np.array([[0., 1., 1.5], [7., 8., -5.5]]).astype(np.float32)
147 neg, pos = sess.run(None, {'x': x})
148 diff = x - (neg + pos)
149 assert_almost_equal(diff, np.zeros(diff.shape))
150
151 def test_cc_negpos(self):
152 so = _ort.SessionOptions()
153 so.register_custom_ops_library(_get_library_path())
154 onnx_model = _create_test_model_2outputs("")
155 self.assertIn('op_type: "NegPos"', str(onnx_model))
156 sess = _ort.InferenceSession(onnx_model.SerializeToString(), so)
157 x = np.array([[0., 1., 1.5], [7., 8., -5.5]]).astype(np.float32)
158 neg, pos = sess.run(None, {'x': x})
159 diff = x - (neg + pos)
160 assert_almost_equal(diff, np.zeros(diff.shape))
161
162 def test_check_saved_model(self):
163 this = os.path.dirname(__file__)
164 so = _ort.SessionOptions()
165 so.register_custom_ops_library(_get_library_path())
166 onnx_content = _create_test_model_test()
167 onnx_bytes = onnx_content.SerializeToString()
168 with open(os.path.join(this, 'data', 'custom_op_test.onnx'),
169 'rb') as f:
170 saved = f.read()
171 assert onnx_bytes == saved
172
173 def test_cc_operator(self):
174 so = _ort.SessionOptions()
175 so.register_custom_ops_library(_get_library_path())
176 onnx_content = _create_test_model_test()
177 self.assertIn('op_type: "CustomOpOne"', str(onnx_content))
178 ser = onnx_content.SerializeToString()
179 sess0 = _ort.InferenceSession(ser, so)
180 res = sess0.run(None, {
181 'input_1': np.random.rand(3, 5).astype(np.float32),
182 'input_2': np.random.rand(3, 5).astype(np.float32)})
183 self.assertEqual(res[0].shape, (3, 5))
184
185
186if __name__ == "__main__":
187 unittest.main()
188