add ort_optimized_model

Author: sdpython

.gitignore

@@ -10,3 +10,5 @@ build/*
 *egg-info/*
 _doc/auto_examples/*
 _doc/examples/plot_*.png
+_doc/_static/require.js
+_doc/_static/viz.js

_doc/examples/plot_onnxruntime.py

@@ -56,11 +56,21 @@ def myloss(x, y):
 yort = OrtTensor.from_array(y)
 
 
-def loop():
-    for _ in range(1000):
+def loop_ort(n):
+    for _ in range(n):
         ort_myloss(xort, yort)
 
 
+def loop_numpy(n):
+    for _ in range(n):
+        myloss(x, y)
+
+
+def loop(n=1000):
+    loop_numpy(n)
+    loop_ort(n)
+
+
 ps = profile(loop)[0]
 root, nodes = profile2graph(ps, clean_text=lambda x: x.split("/")[-1])
 text = root.to_text()

_unittests/ut_ort/test_ort_opimizer.py

@@ -0,0 +1,36 @@
+import unittest
+
+import numpy as np
+from onnx.defs import onnx_opset_version
+from onnx_array_api.npx import absolute, jit_onnx
+from onnx_array_api.ext_test_case import ExtTestCase
+from onnx_array_api.ort.ort_optimizers import ort_optimized_model
+
+
+DEFAULT_OPSET = onnx_opset_version()
+
+
+class TestOrtOptimizer(ExtTestCase):
+    def test_ort_optimizers(self):
+        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()
+        self.assertRaise(lambda: ort_optimized_model(onx, "NO"), ValueError)
+        optimized = ort_optimized_model(onx)
+        self.assertIn('op_type: "Squeeze"', str(optimized))
+        self.assertIn("initializer {", str(optimized))
+
+
+if __name__ == "__main__":
+    unittest.main(verbosity=2)

_unittests/ut_ort/test_ort_tensor.py

@@ -3,7 +3,6 @@
 from io import StringIO
 
 import numpy as np
-import scipy
 from onnx.defs import onnx_opset_version
 from onnx.reference import ReferenceEvaluator
 from onnxruntime import InferenceSession
@@ -21,7 +20,6 @@
 
 
 class TestOrtTensor(ExtTestCase):
-    @unittest.skipIf(InferenceSession is None, reason="onnxruntime is needed.")
     def test_eager_numpy_type_ort(self):
         def impl(A):
             self.assertIsInstance(A, EagerOrtTensor)
@@ -47,7 +45,6 @@ def impl(A):
         self.assertEqualArray(z, res.numpy())
         self.assertEqual(res.numpy().dtype, np.float64)
 
-    @unittest.skipIf(InferenceSession is None, reason="onnxruntime is needed.")
     def test_eager_numpy_type_ort_op(self):
         def impl(A):
             self.assertIsInstance(A, EagerOrtTensor)
@@ -71,7 +68,6 @@ def impl(A):
         self.assertEqualArray(z, res.numpy())
         self.assertEqual(res.numpy().dtype, np.float64)
 
-    @unittest.skipIf(InferenceSession is None, reason="onnxruntime is needed.")
     def test_eager_ort(self):
         def impl(A):
             print("A")
@@ -145,8 +141,6 @@ def impl(A):
         self.assertEqual(tuple(res.shape()), z.shape)
         self.assertStartsWith("A\nB\nC\n", text)
 
-    @unittest.skipIf(InferenceSession is None, reason="onnxruntime is not available")
-    @unittest.skipIf(scipy is None, reason="scipy is not installed.")
     def test_cdist_com_microsoft(self):
         from scipy.spatial.distance import cdist as scipy_cdist
 

onnx_array_api/cache.py

@@ -0,0 +1,19 @@
+from pathlib import Path
+
+
+def get_cache_file(filename: str, remove: bool = False):
+    """
+    Returns a name in the cache folder `~/.onnx-array-api`.
+
+    :param filename: filename
+    :param remove: remove if exists
+    :return: full filename
+    """
+    home = Path.home()
+    folder = home / ".onnx-array-api"
+    if not folder.exists():
+        folder.mkdir()
+    name = folder / filename
+    if name.exists():
+        name.unlink()
+    return name

onnx_array_api/ort/ort_optimizers.py

@@ -0,0 +1,37 @@
+from typing import Union
+from onnx import ModelProto, load
+from onnxruntime import InferenceSession, SessionOptions
+from onnxruntime.capi._pybind_state import GraphOptimizationLevel
+from ..cache import get_cache_file
+
+
+def ort_optimized_model(
+    onx: Union[str, ModelProto], level: str = "ORT_ENABLE_ALL"
+) -> Union[str, ModelProto]:
+    """
+    Returns the optimized model used by onnxruntime before
+    running computing the inference.
+
+    :param onx: ModelProto
+    :param level: optimization level, `'ORT_ENABLE_BASIC'`,
+        `'ORT_ENABLE_EXTENDED'`, `'ORT_ENABLE_ALL'`
+    :return: optimized model
+    """
+    glevel = getattr(GraphOptimizationLevel, level, None)
+    if glevel is None:
+        raise ValueError(
+            f"Unrecognized level {level!r} among {dir(GraphOptimizationLevel)}."
+        )
+
+    cache = get_cache_file("ort_optimized_model.onnx", remove=True)
+    so = SessionOptions()
+    so.graph_optimization_level = glevel
+    so.optimized_model_filepath = str(cache)
+    InferenceSession(onx if isinstance(onx, str) else onx.SerializeToString(), so)
+    if not cache.exists():
+        raise RuntimeError(f"The optimized model {str(cache)!r} not found.")
+    if isinstance(onx, str):
+        return str(cache)
+    opt_onx = load(str(cache))
+    cache.unlink()
+    return opt_onx

onnx_array_api/ort/ort_tensors.py

@@ -1,7 +1,3 @@
-"""
-This is an example of a backend for classes :class:`JitOnnx` and :class:`JitEager`
-using onnxruntime as a runtime. It is provided as an example.
-"""
 from typing import Any, Callable, List, Optional, Tuple, Union
 
 import numpy as np

onnx_array_api/plotting/dot_plot.py

@@ -254,7 +254,7 @@ def dot_label(text):
                 )
             static_inputs.append(out)
 
-        if node.name.strip() == "":
+        if node.name.strip() == "" or node.name in fill_names:
             name = node.op_type
             iname = 1
             while name in fill_names:
@@ -329,7 +329,7 @@ def dot_label(text):
             exp.append(
                 f"  {dot_name(prefix)}{dot_name(node.name)} "
                 f'[shape=box style="filled,rounded" color=orange '
-                f'label="{node.op_type}\\n({dot_name(node.name)}){satts}" '
+                f'label="{node.op_type}{satts}" '
                 f"fontsize={fontsize}];"
             )
 

setup.py

@@ -24,8 +24,8 @@
     requirements = ["numpy", "scipy", "onnx"]
 
 try:
-    with open(os.path.join(here, "README.rst"), "r", encoding='utf-8') as f:
-        long_description = "onnx-array-api:" + f.read().split('onnx-array-api:')[1]
+    with open(os.path.join(here, "README.rst"), "r", encoding="utf-8") as f:
+        long_description = "onnx-array-api:" + f.read().split("onnx-array-api:")[1]
 except FileNotFoundError:
     long_description = ""