def test_custom_pythonop_pytorch(self): # register_custom_op_symbolic( # '<namespace>::inverse', my_inverse, <opset_version>) register_custom_op_symbolic('::inverse', my_inverse, 1) x = torch.randn(3, 3) # Export model to ONNX f = io.BytesIO() torch.onnx.export(CustomInverse(), (x, ), f) onnx_model = load(io.BytesIO(f.getvalue())) self.assertIn('domain: "ai.onnx.contrib"', str(onnx_model)) model = CustomInverse() pt_outputs = model(x) so = _ort.SessionOptions() so.register_custom_ops_library(_get_library_path()) # Run the exported model with ONNX Runtime ort_sess = _ort.InferenceSession(f.getvalue(), so) 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)
def test_string_sentencepiece_tokenizer_bin(self): so = _ort.SessionOptions() so.register_custom_ops_library(_get_library_path()) model, model_b64 = load_piece('model__6') modelb = bytes(model) py_onnx_model = _create_test_model_sentencepiece('Py', None) self.assertIn('op_type: "PySentencepieceTokenizer"', str(py_onnx_model)) cc_onnx_model = _create_test_model_sentencepiece('', modelb) self.assertIn('op_type: "SentencepieceTokenizer"', str(cc_onnx_model)) py_sess = _ort.InferenceSession(py_onnx_model.SerializeToString(), so) cc_sess = _ort.InferenceSession(cc_onnx_model.SerializeToString(), so) alpha = 0 nbest_size = 0 bools = 0 inputs = dict(model=model, inputs=np.array(["Hello world", "Hello world louder"], dtype=np.object), nbest_size=np.array([nbest_size], dtype=np.int64), alpha=np.array([alpha], dtype=np.float32), add_bos=np.array([bools & 1], dtype=np.bool_), add_eos=np.array([bools & 2], dtype=np.bool_), reverse=np.array([bools & 4], dtype=np.bool_)) exp = self.SentencepieceTokenizer(**inputs) py_txout = py_sess.run(None, inputs) del inputs['model'] cc_txout = cc_sess.run(None, inputs) for i in range(0, 2): assert_almost_equal(exp[i], py_txout[i]) assert_almost_equal(exp[i], cc_txout[i])
def test_string_ragged_string_to_dense_cc(self): so = _ort.SessionOptions() so.register_custom_ops_library(_get_library_path()) model, model_b64 = load_piece('model__6') onnx_model = _create_test_model_ragged_to_dense('', model_b64) self.assertIn('op_type: "RaggedTensorToDense"', str(onnx_model)) sess = _ort.InferenceSession(onnx_model.SerializeToString(), so) inputs = dict(model=model, inputs=np.array(["Hello world", "Hello world louder"], dtype=np.object), nbest_size=np.array([0], dtype=np.int64), alpha=np.array([0], dtype=np.float32), add_bos=np.array([0], dtype=np.bool_), add_eos=np.array([0], dtype=np.bool_), reverse=np.array([0], dtype=np.bool_)) del inputs['model'] txout = sess.run(None, inputs) assert_almost_equal( txout[0], np.array([[17486, 1017, -1, -1], [17486, 1017, 155, 21869]], dtype=np.int64)) assert_almost_equal( txout[1], np.array([17486, 1017, 17486, 1017, 155, 21869], dtype=np.int32))
def test_check_saved_model(self): this = os.path.dirname(__file__) so = _ort.SessionOptions() so.register_custom_ops_library(_get_library_path()) onnx_content = _create_test_model_test() onnx_bytes = onnx_content.SerializeToString() with open(os.path.join(this, 'data', 'custom_op_test.onnx'), 'rb') as f: saved = f.read() assert onnx_bytes == saved
def test_python_join(self): so = _ort.SessionOptions() so.register_custom_ops_library(_get_library_path()) onnx_model = _create_test_join() self.assertIn('op_type: "PyOpJoin"', str(onnx_model)) sess = _ort.InferenceSession(onnx_model.SerializeToString(), so) arr = np.array([["a", "b"]], dtype=np.object) txout = sess.run(None, {'input_1': arr}) exp = np.array(["a;b"], dtype=np.object) assert txout[0][0] == exp[0]
def test_cc_negpos(self): so = _ort.SessionOptions() so.register_custom_ops_library(_get_library_path()) onnx_model = _create_test_model_2outputs("") self.assertIn('op_type: "NegPos"', str(onnx_model)) sess = _ort.InferenceSession(onnx_model.SerializeToString(), so) x = np.array([[0., 1., 1.5], [7., 8., -5.5]]).astype(np.float32) neg, pos = sess.run(None, {'x': x}) diff = x - (neg + pos) assert_almost_equal(diff, np.zeros(diff.shape))
def test_add_epsilon_python(self): so = _ort.SessionOptions() so.register_custom_ops_library(_get_library_path()) onnx_model = _create_test_model_double('Py') self.assertIn('op_type: "PyAddEpsilon"', str(onnx_model)) sess = _ort.InferenceSession(onnx_model.SerializeToString(), so) input_1 = np.array([[0., 1., 1.5], [7., 8., -5.5]]) txout = sess.run(None, {'input_1': input_1}) diff = txout[0] - input_1 - 1e-3 assert_almost_equal(diff, np.zeros(diff.shape))
def test_python_operator(self): so = _ort.SessionOptions() so.register_custom_ops_library(_get_library_path()) onnx_model = _create_test_model() self.assertIn('op_type: "PyReverseMatrix"', str(onnx_model)) sess = _ort.InferenceSession(onnx_model.SerializeToString(), so) input_1 = np.array([1, 2, 3, 4, 5, 6]).astype(np.float32).reshape([3, 2]) txout = sess.run(None, {'input_1': input_1}) assert_almost_equal(txout[0], np.array([[5., 6.], [3., 4.], [1., 2.]]))
def _run_string_length(input): model = _create_test_model(input.ndim, input.ndim) so = _ort.SessionOptions() so.register_custom_ops_library(_get_library_path()) sess = _ort.InferenceSession(model.SerializeToString(), so) result = sess.run(None, {'input_1': input}) # verify output = np.array([len(elem) for elem in input.flatten()], dtype=np.int64).reshape(input.shape) np.testing.assert_array_equal(result, [output])
def test_cc_operator(self): so = _ort.SessionOptions() so.register_custom_ops_library(_get_library_path()) onnx_content = _create_test_model_test() self.assertIn('op_type: "CustomOpOne"', str(onnx_content)) ser = onnx_content.SerializeToString() sess0 = _ort.InferenceSession(ser, so) res = sess0.run( None, { 'input_1': np.random.rand(3, 5).astype(np.float32), 'input_2': np.random.rand(3, 5).astype(np.float32) }) self.assertEqual(res[0].shape, (3, 5))
def _run_tokenizer(self, test_sentence, padding_length=-1): model = _create_test_model(vocab_file=self.tokjson, merges_file=self.merges, max_length=padding_length) so = _ort.SessionOptions() so.register_custom_ops_library(_get_library_path()) sess = _ort.InferenceSession(model.SerializeToString(), so) input_text = np.array(test_sentence) input_ids, attention_mask = sess.run(None, {'string_input': input_text}) expect_input_ids, expect_attention_mask = self.tokenizer.tokenizer_sentence( test_sentence, padding_length) np.testing.assert_array_equal(expect_input_ids, input_ids) np.testing.assert_array_equal(expect_attention_mask, attention_mask) del sess del so
def _run_string_concat(input1, input2): model = _create_test_model(input1.ndim, input1.ndim) so = _ort.SessionOptions() so.register_custom_ops_library(_get_library_path()) sess = _ort.InferenceSession(model.SerializeToString(), so) result = sess.run(None, {'input_1': input1, 'input_2': input2}) # verify output = [] shape = input1.shape input1 = input1.flatten() input2 = input2.flatten() for i in range(len(input1)): output.append(input1[i] + input2[i]) output = np.array(output).reshape(shape) np.testing.assert_array_equal(result, [output])
def test_pyop_hooking(self): # type: () -> None model = torchvision.models.mobilenet_v2(pretrained=False) x = torch.rand(1, 3, 224, 224) with io.BytesIO() as f: torch.onnx.export(model, (x, ), f) model = onnx.load_model_from_string(f.getvalue()) self.assertTrue(model.graph.node[5].op_type == 'Conv') hkd_model = hook_model_op(model, model.graph.node[5].name, TestPyTorchCustomOp.on_hook, [PyOp.dt_float] * 3) so = _ort.SessionOptions() so.register_custom_ops_library(_get_library_path()) sess = _ort.InferenceSession(hkd_model.SerializeToString(), so) TestPyTorchCustomOp._hooked = False sess.run(None, {'input.1': x.numpy()}) self.assertTrue(TestPyTorchCustomOp._hooked)
def test_string_ragged_string_to_sparse_python(self): so = _ort.SessionOptions() so.register_custom_ops_library(_get_library_path()) model, model_b64 = load_piece('model__6') onnx_model = _create_test_model_ragged_to_sparse('Py', None) self.assertIn('op_type: "PyRaggedTensorToSparse"', str(onnx_model)) sess = _ort.InferenceSession(onnx_model.SerializeToString(), so) inputs = dict(model=model, inputs=np.array(["Hello world", "Hello world louder"], dtype=np.object), nbest_size=np.array([0], dtype=np.int64), alpha=np.array([0], dtype=np.float32), add_bos=np.array([0], dtype=np.bool_), add_eos=np.array([0], dtype=np.bool_), reverse=np.array([0], dtype=np.bool_)) txout = sess.run(None, inputs) temp = self.SentencepieceTokenizer(**inputs) exp = self.RaggedTensorToSparse(temp[1], temp[0]) for i in range(0, 3): assert_almost_equal(exp[i], txout[i])