microsoft/onnxruntime-extensions
Publicmirrored fromhttps://github.com/microsoft/onnxruntime-extensionsAvailable
test/test_pyops.py
46lines · modecode
| 1 | import numpy as np |
| 2 | from onnx import helper, onnx_pb as onnx_proto |
| 3 | import onnxruntime as _ort |
| 4 | from ortcustomops import ( |
| 5 | onnx_op, |
| 6 | get_library_path as _get_library_path) |
| 7 | |
| 8 | |
| 9 | def _create_test_model(): |
| 10 | nodes = [] |
| 11 | nodes[0:] = [helper.make_node('Identity', ['input_1'], ['identity1'])] |
| 12 | nodes[1:] = [helper.make_node('ReverseMatrix', |
| 13 | ['identity1'], ['reversed'], |
| 14 | domain='ai.onnx.contrib')] |
| 15 | |
| 16 | input0 = helper.make_tensor_value_info('input_1', onnx_proto.TensorProto.FLOAT, [None, 2]) |
| 17 | output0 = helper.make_tensor_value_info('reversed', onnx_proto.TensorProto.FLOAT, [None, 2]) |
| 18 | |
| 19 | graph = helper.make_graph(nodes, 'test0', [input0], [output0]) |
| 20 | model = helper.make_model(graph, opset_imports=[helper.make_operatorsetid('ai.onnx.contrib', 1)]) |
| 21 | return model |
| 22 | |
| 23 | |
| 24 | @onnx_op(op_type="ReverseMatrix") |
| 25 | def reverse_matrix(x): |
| 26 | # the user custom op implementation here: |
| 27 | return np.flip(x, axis=0) |
| 28 | |
| 29 | |
| 30 | # TODO: refactor the following code into pytest cases, right now, the script is more friendly for debugging. |
| 31 | so = _ort.SessionOptions() |
| 32 | so.register_custom_ops_library(_get_library_path()) |
| 33 | |
| 34 | sess0 = _ort.InferenceSession('./test/data/custom_op_test.onnx', so) |
| 35 | |
| 36 | res = sess0.run(None, { |
| 37 | 'input_1': np.random.rand(3, 5).astype(np.float32), 'input_2': np.random.rand(3, 5).astype(np.float32)}) |
| 38 | |
| 39 | print(res[0]) |
| 40 | |
| 41 | sess = _ort.InferenceSession(_create_test_model().SerializeToString(), so) |
| 42 | |
| 43 | txout = sess.run(None, { |
| 44 | 'input_1': np.array([1, 2, 3, 4, 5, 6]).astype(np.float32).reshape([3, 2])}) |
| 45 | |
| 46 | print(txout[0]) |