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
b5ba84a185c5c85c436e744c40ac9a8cbd9d3f1f

Branches

Tags

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

Clone

HTTPS

Download ZIP

onnxruntime-extensions/test

test/test_pyops.py

238lines · 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
91def _create_test_join():
92 nodes = []
93 nodes[0:] = [helper.make_node('Identity', ['input_1'], ['identity1'])]
94 nodes[1:] = [helper.make_node('PyOpJoin',
95 ['identity1'], ['joined'],
96 sep=';',
97 domain='ai.onnx.contrib')]
98
99 input0 = helper.make_tensor_value_info(
100 'input_1', onnx_proto.TensorProto.STRING, [None, None])
101 output0 = helper.make_tensor_value_info(
102 'joined', onnx_proto.TensorProto.STRING, [None])
103
104 graph = helper.make_graph(nodes, 'test0', [input0], [output0])
105 model = helper.make_model(
106 graph, opset_imports=[helper.make_operatorsetid('ai.onnx.contrib', 1)])
107 return model
108
109
110class TestPythonOp(unittest.TestCase):
111
112 @classmethod
113 def setUpClass(cls):
114
115 @onnx_op(op_type="CustomOpOne",
116 inputs=[PyCustomOpDef.dt_float, PyCustomOpDef.dt_float])
117 def custom_one_op(x, y):
118 return np.add(x, y)
119
120 @onnx_op(op_type="CustomOpTwo",
121 outputs=[PyCustomOpDef.dt_int32])
122 def custom_two_op(f):
123 return np.round(f).astype(np.int32)
124
125 @onnx_op(op_type="PyReverseMatrix")
126 def reverse_matrix(x):
127 # The user custom op implementation here.
128 return np.flip(x, axis=0).astype(np.float32)
129
130 @onnx_op(op_type="PyAddEpsilon",
131 inputs=[PyCustomOpDef.dt_double],
132 outputs=[PyCustomOpDef.dt_double])
133 def add_epsilon(x):
134 # The user custom op implementation here.
135 return x + 1e-3
136
137 @onnx_op(op_type="PyNegPos",
138 inputs=[PyCustomOpDef.dt_float],
139 outputs=[PyCustomOpDef.dt_float, PyCustomOpDef.dt_float])
140 def negpos(x):
141 neg = x.copy()
142 pos = x.copy()
143 neg[x > 0] = 0
144 pos[x < 0] = 0
145 return neg, pos
146
147 @onnx_op(op_type="PyOpJoin",
148 inputs=[PyCustomOpDef.dt_string],
149 outputs=[PyCustomOpDef.dt_string],
150 attrs=['sep'])
151 def join(xs, **kwargs):
152 sep = kwargs.get('sep', '')
153 res = []
154 for x in xs:
155 res.append(sep.join(x))
156 return np.array(res, dtype=np.object)
157
158 def test_python_operator(self):
159 so = _ort.SessionOptions()
160 so.register_custom_ops_library(_get_library_path())
161 onnx_model = _create_test_model()
162 self.assertIn('op_type: "PyReverseMatrix"', str(onnx_model))
163 sess = _ort.InferenceSession(onnx_model.SerializeToString(), so)
164 input_1 = np.array(
165 [1, 2, 3, 4, 5, 6]).astype(np.float32).reshape([3, 2])
166 txout = sess.run(None, {'input_1': input_1})
167 assert_almost_equal(txout[0], np.array([[5., 6.], [3., 4.], [1., 2.]]))
168
169 def test_add_epsilon_python(self):
170 so = _ort.SessionOptions()
171 so.register_custom_ops_library(_get_library_path())
172 onnx_model = _create_test_model_double('Py')
173 self.assertIn('op_type: "PyAddEpsilon"', str(onnx_model))
174 sess = _ort.InferenceSession(onnx_model.SerializeToString(), so)
175 input_1 = np.array([[0., 1., 1.5], [7., 8., -5.5]])
176 txout = sess.run(None, {'input_1': input_1})
177 diff = txout[0] - input_1 - 1e-3
178 assert_almost_equal(diff, np.zeros(diff.shape))
179
180 def test_python_negpos(self):
181 so = _ort.SessionOptions()
182 so.register_custom_ops_library(_get_library_path())
183 onnx_model = _create_test_model_2outputs('Py')
184 self.assertIn('op_type: "PyNegPos"', str(onnx_model))
185 sess = _ort.InferenceSession(onnx_model.SerializeToString(), so)
186 x = np.array([[0., 1., 1.5], [7., 8., -5.5]]).astype(np.float32)
187 neg, pos = sess.run(None, {'x': x})
188 diff = x - (neg + pos)
189 assert_almost_equal(diff, np.zeros(diff.shape))
190
191 def test_cc_negpos(self):
192 so = _ort.SessionOptions()
193 so.register_custom_ops_library(_get_library_path())
194 onnx_model = _create_test_model_2outputs("")
195 self.assertIn('op_type: "NegPos"', str(onnx_model))
196 sess = _ort.InferenceSession(onnx_model.SerializeToString(), so)
197 x = np.array([[0., 1., 1.5], [7., 8., -5.5]]).astype(np.float32)
198 neg, pos = sess.run(None, {'x': x})
199 diff = x - (neg + pos)
200 assert_almost_equal(diff, np.zeros(diff.shape))
201
202 def test_check_saved_model(self):
203 this = os.path.dirname(__file__)
204 so = _ort.SessionOptions()
205 so.register_custom_ops_library(_get_library_path())
206 onnx_content = _create_test_model_test()
207 onnx_bytes = onnx_content.SerializeToString()
208 with open(os.path.join(this, 'data', 'custom_op_test.onnx'),
209 'rb') as f:
210 saved = f.read()
211 assert onnx_bytes == saved
212
213 def test_cc_operator(self):
214 so = _ort.SessionOptions()
215 so.register_custom_ops_library(_get_library_path())
216 onnx_content = _create_test_model_test()
217 self.assertIn('op_type: "CustomOpOne"', str(onnx_content))
218 ser = onnx_content.SerializeToString()
219 sess0 = _ort.InferenceSession(ser, so)
220 res = sess0.run(None, {
221 'input_1': np.random.rand(3, 5).astype(np.float32),
222 'input_2': np.random.rand(3, 5).astype(np.float32)})
223 self.assertEqual(res[0].shape, (3, 5))
224
225 def test_python_join(self):
226 so = _ort.SessionOptions()
227 so.register_custom_ops_library(_get_library_path())
228 onnx_model = _create_test_join()
229 self.assertIn('op_type: "PyOpJoin"', str(onnx_model))
230 sess = _ort.InferenceSession(onnx_model.SerializeToString(), so)
231 arr = np.array([["a", "b"]], dtype=np.object)
232 txout = sess.run(None, {'input_1': arr})
233 exp = np.array(["a;b"], dtype=np.object)
234 assert txout[0][0] == exp[0]
235
236
237if __name__ == "__main__":
238 unittest.main()
239