Supports OrtValue in function ort_profile

sdpython · sdpython · Jun 22, 2023 · Jun 22, 2023 · Jun 22, 2023 · Jun 22, 2023
commit a763dc0f798642f7bd7499477942d7f652b754ad
diff --git a/.gitignore b/.gitignore
@@ -10,8 +10,10 @@ build/*
 .eggs/*
 .hypothesis/*
 *egg-info/*
+onnxruntime_profile*
 _doc/auto_examples/*
 _doc/examples/_cache/*
+_doc/examples/onnxruntime_profile*
 _doc/examples/plot_*.png
 _doc/examples/plot_*.xlsx
 _doc/examples/data/*.optimized.onnx

diff --git a/CHANGELOGS.rst b/CHANGELOGS.rst
@@ -4,5 +4,6 @@ Change Logs
 0.2.0
 +++++
 
+* :pr:`27`: support OrtValue in function :func:`ort_profile`
 * :pr:`17`: implements ArrayAPI
 * :pr:`3`: fixes Array API with onnxruntime and scikit-learn
diff --git a/_unittests/ut_ort/test_ort_profile.py b/_unittests/ut_ort/test_ort_profile.py
@@ -6,6 +6,10 @@
 from onnx_array_api.ext_test_case import ExtTestCase
 from onnx_array_api.ort.ort_optimizers import ort_optimized_model
 from onnx_array_api.ort.ort_profile import ort_profile, merge_ort_profile
+from onnxruntime.capi._pybind_state import (
+    OrtValue as C_OrtValue,
+    OrtDevice as C_OrtDevice,
+)
 
 
 class TestOrtProfile(ExtTestCase):
@@ -28,7 +32,36 @@ def myloss(x, y):
         self.assertRaise(lambda: ort_optimized_model(onx, "NO"), ValueError)
         optimized = ort_optimized_model(onx)
         prof = ort_profile(optimized, feeds)
-        prof.to_csv("prof.csv", index=False)
+        self.assertIsInstance(prof, DataFrame)
+        prof = ort_profile(optimized, feeds, as_df=False)
+        self.assertIsInstance(prof, list)
+
+    def test_ort_profile_ort_value(self):
+        def to_ort_value(m):
+            device = C_OrtDevice(C_OrtDevice.cpu(), C_OrtDevice.default_memory(), 0)
+            ort_value = C_OrtValue.ortvalue_from_numpy(m, device)
+            return ort_value
+
+        def l1_loss(x, y):
+            return absolute(x - y).sum()
+
+        def l2_loss(x, y):
+            return ((x - y) ** 2).sum()
+
+        def myloss(x, y):
+            return l1_loss(x[:, 0], y[:, 0]) + l2_loss(x[:, 1], y[:, 1])
+
+        jitted_myloss = jit_onnx(myloss)
+        x = np.array([[0.1, 0.2], [0.3, 0.4]], dtype=np.float32)
+        y = np.array([[0.11, 0.22], [0.33, 0.44]], dtype=np.float32)
+        jitted_myloss(x, y)
+        onx = jitted_myloss.get_onnx()
+        np_feeds = {"x0": x, "x1": y}
+        feeds = {k: to_ort_value(v) for k, v in np_feeds.items()}
+
+        self.assertRaise(lambda: ort_optimized_model(onx, "NO"), ValueError)
+        optimized = ort_optimized_model(onx)
+        prof = ort_profile(optimized, feeds)
         self.assertIsInstance(prof, DataFrame)
         prof = ort_profile(optimized, feeds, as_df=False)
         self.assertIsInstance(prof, list)

diff --git a/onnx_array_api/ort/ort_profile.py b/onnx_array_api/ort/ort_profile.py
@@ -45,8 +45,16 @@ def ort_profile(
     if providers is None:
         providers = ["CPUExecutionProvider"]
     sess = InferenceSession(obj, sess_options, providers=providers, **kwargs)
-    for i in range(repeat):
-        sess.run(None, feeds)
+    first = list(feeds.values())[0]
+
+    if isinstance(first, numpy.ndarray):
+        for i in range(repeat):
+            sess.run(None, feeds)
+    else:
+        out_names = [o.name for o in sess.get_outputs()]
+        for i in range(repeat):
+            sess._sess.run_with_ort_values(feeds, out_names, None)
+
     prof = sess.end_profiling()
     with open(prof, "r") as f:
         content = f.read()