Add onnx support (#292)

maheshambule · aguschin · web-flow · commit bed41adcbde3 · 2022-07-05T15:55:24.000+06:00
* onnx support first cut

* incorporated review comments

* add backport, move skl2onnx to test dependency, store onnx model tensors as external data

* fix setup.py tests

* fix setup.py tests

* update dependencies

* update dependencies

* fix flaky testcase

* Update tests/contrib/test_pandas.py

Co-authored-by: Alexander Guschin &lt;1aguschin@gmail.com&gt;
diff --git a/mlem/contrib/onnx.py b/mlem/contrib/onnx.py
@@ -0,0 +1,139 @@
+from typing import Any, ClassVar, List, Optional, Union
+
+import numpy as np
+import onnx
+import onnxruntime as onnxrt
+import pandas as pd
+from numpy.typing import DTypeLike
+from onnx import ModelProto, ValueInfoProto
+from onnx.mapping import TENSOR_TYPE_TO_NP_TYPE
+
+from mlem.core.artifacts import Artifacts, Storage
+from mlem.core.hooks import IsInstanceHookMixin
+from mlem.core.model import ModelHook, ModelIO, ModelType, Signature
+from mlem.core.requirements import InstallableRequirement, Requirements
+from mlem.utils.backport import cached_property
+from mlem.utils.module import get_object_requirements
+
+
+def convert_to_numpy(
+    data: Union[np.ndarray, pd.DataFrame], dtype: DTypeLike
+) -> np.ndarray:
+    """Converts input data to numpy"""
+    if isinstance(data, np.ndarray):
+        pass
+    elif isinstance(data, pd.DataFrame):
+        data = data.to_numpy()
+    else:
+        raise TypeError(f"input data type: {type(data)} is not supported")
+    return data.astype(dtype=dtype)
+
+
+def get_onnx_to_numpy_type(value_info: ValueInfoProto) -> DTypeLike:
+    """Returns numpy equivalent type of onnx value info"""
+    onnx_type = value_info.type.tensor_type.elem_type
+    return TENSOR_TYPE_TO_NP_TYPE[onnx_type]
+
+
+class ModelProtoIO(ModelIO):
+    """IO for ONNX model object"""
+
+    type: ClassVar[str] = "model_proto"
+
+    def dump(self, storage: Storage, path: str, model) -> Artifacts:
+        path = f"{path}/model.onnx"
+        with storage.open(path) as (f, art):
+            onnx.save_model(
+                model,
+                f,
+                save_as_external_data=True,
+                location="tensors",
+                size_threshold=0,
+                all_tensors_to_one_file=True,
+            )
+        return {self.art_name: art}
+
+    def load(self, artifacts: Artifacts):
+        if len(artifacts) != 1:
+            raise ValueError("Invalid artifacts: should be one .onnx file")
+        with artifacts[self.art_name].open() as f:
+            return onnx.load_model(f)
+
+
+class ONNXModel(ModelType, ModelHook, IsInstanceHookMixin):
+    """
+    :class:`mlem.core.model.ModelType` implementation for `onnx` models
+    """
+
+    type: ClassVar[str] = "onnx"
+    io: ModelIO = ModelProtoIO()
+    valid_types: ClassVar = (ModelProto,)
+
+    class Config:
+        keep_untouched = (cached_property,)
+
+    @classmethod
+    def process(
+        cls, obj: Any, sample_data: Optional[Any] = None, **kwargs
+    ) -> ModelType:
+
+        model = ONNXModel(io=ModelProtoIO(), methods={}).bind(obj)
+        # TODO - use ONNX infer shapes.
+        onnxrt_predict = Signature.from_method(
+            model.predict, auto_infer=sample_data is not None, data=sample_data
+        )
+        model.methods = {
+            "predict": onnxrt_predict,
+        }
+
+        return model
+
+    @cached_property
+    def runtime_session(self) -> onnxrt.InferenceSession:
+        """Provides onnx runtime inference session"""
+        # TODO - add support for runtime providers, options. add support for GPU devices.
+        return onnxrt.InferenceSession(self.model.SerializeToString())
+
+    def predict(self, data: Union[List, np.ndarray, pd.DataFrame]) -> Any:
+        """Returns inference output for given input data"""
+        model_inputs = self.runtime_session.get_inputs()
+
+        if not isinstance(data, list):
+            data = [data]
+
+        if len(model_inputs) != len(data):
+            raise ValueError(
+                f"no of inputs provided: {len(data)}, "
+                f"expected: {len(model_inputs)}"
+            )
+
+        input_dict = {}
+        for model_input, input_data in zip(self.model.graph.input, data):
+            input_dict[model_input.name] = convert_to_numpy(
+                input_data, get_onnx_to_numpy_type(model_input)
+            )
+
+        label_names = [out.name for out in self.runtime_session.get_outputs()]
+        pred_onnx = self.runtime_session.run(label_names, input_dict)
+
+        output = []
+        for output_data in pred_onnx:
+            if isinstance(
+                output_data, list
+            ):  # TODO - temporary workaround to fix fastapi model issues
+                output.append(pd.DataFrame(output_data).to_numpy())
+            else:
+                output.append(output_data)
+
+        return output
+
+    def get_requirements(self) -> Requirements:
+        return (
+            super().get_requirements()
+            + InstallableRequirement.from_module(onnx)
+            + get_object_requirements(self.predict)
+            + Requirements.new(
+                InstallableRequirement(module="protobuf", version="3.20.1")
+            )
+        )
+        # https://github.com/protocolbuffers/protobuf/issues/10051
diff --git a/mlem/core/objects.py b/mlem/core/objects.py
@@ -609,7 +609,9 @@ def from_obj(
         params: Dict[str, str] = None,
     ) -> "MlemModel":
         mt = ModelAnalyzer.analyze(model, sample_data=sample_data)
-        mt.model = model
+        if mt.model is None:
+            mt = mt.bind(model)
+
         return MlemModel(
             model_type=mt,
             requirements=mt.get_requirements().expanded,
diff --git a/mlem/ext.py b/mlem/ext.py
@@ -92,6 +92,7 @@ class ExtensionLoader:
         Extension("mlem.contrib.numpy", ["numpy"], False),
         Extension("mlem.contrib.pandas", ["pandas"], False),
         Extension("mlem.contrib.sklearn", ["sklearn"], False),
+        Extension("mlem.contrib.onnx", ["onnx"], False),
         Extension("mlem.contrib.tensorflow", ["tensorflow"], False),
         Extension("mlem.contrib.torch", ["torch"], False),
         Extension("mlem.contrib.catboost", ["catboost"], False),
diff --git a/mlem/utils/backport.py b/mlem/utils/backport.py
@@ -0,0 +1,59 @@
+import functools
+import sys
+
+if sys.version_info >= (3, 8):
+    cached_property = functools.cached_property
+else:
+    # Code copied from Python 3.8 https://github.com/python/cpython/blob/3.8/Lib/functools.py
+    # cached_property is not available in Python versions < 3.8.
+    from _thread import RLock
+
+    _NOT_FOUND = object()
+
+    class cached_property:
+        def __init__(self, func):
+            self.func = func
+            self.attrname = None
+            self.__doc__ = func.__doc__
+            self.lock = RLock()
+
+        def __set_name__(self, owner, name):
+            if self.attrname is None:
+                self.attrname = name
+            elif name != self.attrname:
+                raise TypeError(
+                    "Cannot assign the same cached_property to two different names "
+                    f"({self.attrname!r} and {name!r})."
+                )
+
+        def __get__(self, instance, owner=None):
+            if instance is None:
+                return self
+            if self.attrname is None:
+                raise TypeError(
+                    "Cannot use cached_property instance without calling __set_name__ on it."
+                )
+            try:
+                cache = instance.__dict__
+            except AttributeError:  # not all objects have __dict__ (e.g. class defines slots)
+                msg = (
+                    f"No '__dict__' attribute on {type(instance).__name__!r} "
+                    f"instance to cache {self.attrname!r} property."
+                )
+                raise TypeError(msg) from None
+            val = cache.get(self.attrname, _NOT_FOUND)
+            if val is _NOT_FOUND:
+                with self.lock:
+                    # check if another thread filled cache while we awaited lock
+                    val = cache.get(self.attrname, _NOT_FOUND)
+                    if val is _NOT_FOUND:
+                        val = self.func(instance)
+                        try:
+                            cache[self.attrname] = val
+                        except TypeError:
+                            msg = (
+                                f"The '__dict__' attribute on {type(instance).__name__!r} instance "
+                                f"does not support item assignment for caching {self.attrname!r} property."
+                            )
+                            raise TypeError(msg) from None
+            return val
diff --git a/mlem/utils/module.py b/mlem/utils/module.py
@@ -343,7 +343,9 @@ def lstrip_lines(lines: Union[str, List[str]], check=True) -> str:
 
 
 _SKIP_CLOSURE_OBJECTS: Dict[str, Dict[str, Set[str]]] = {
-    "globals": {"re": {"_cache"}},
+    # In onnx, "protobuf" module is imported using "google.protobuf" namespace which results in identifying "google"
+    # as possible installable requirement which is incorrect. TODO - see if this can be handled in more correct way
+    "globals": {"re": {"_cache"}, "onnx": {"google"}},
     "nonlocals": {},
 }
 
diff --git a/setup.py b/setup.py
@@ -46,6 +46,7 @@
     "xlrd",
     "tables",
     "pyarrow",
+    "skl2onnx",
 ]
 
 extras = {
@@ -54,6 +55,11 @@
     "pandas": ["pandas"],
     "numpy": ["numpy"],
     "sklearn": ["scikit-learn"],
+    "onnx": ["onnx"],
+    "onnxruntime": [
+        "protobuf==3.20.0",
+        "onnxruntime",
+    ],  # TODO - see if it can be merged with onnx
     "catboost": ["catboost"],
     "xgboost": ["xgboost"],
     "lightgbm": ["lightgbm"],
@@ -151,13 +157,15 @@
             "model_io.lightgbm_io = mlem.contrib.lightgbm:LightGBMModelIO",
             "model_io.pickle = mlem.contrib.callable:PickleModelIO",
             "model_io.xgboost_io = mlem.contrib.xgboost:XGBoostModelIO",
+            "model_io.model_proto = mlem.contrib.onnx:ModelProtoIO",
             "model_io.torch_io = mlem.contrib.torch:TorchModelIO",
             "model_io.tf_keras = mlem.contrib.tensorflow:TFKerasModelIO",
             "model_type.callable = mlem.contrib.callable:CallableModelType",
             "model_type.catboost = mlem.contrib.catboost:CatBoostModel",
             "model_type.lightgbm = mlem.contrib.lightgbm:LightGBMModel",
             "model_type.sklearn = mlem.contrib.sklearn:SklearnModel",
             "model_type.sklearn_pipeline = mlem.contrib.sklearn:SklearnPipelineType",
+            "model_type.onnx = mlem.contrib.onnx:ONNXModel",
             "model_type.xgboost = mlem.contrib.xgboost:XGBoostModel",
             "model_type.torch = mlem.contrib.torch:TorchModel",
             "model_type.tf_keras = mlem.contrib.tensorflow:TFKerasModel",
diff --git a/tests/contrib/test_onnx.py b/tests/contrib/test_onnx.py
diff --git a/tests/contrib/test_pandas.py b/tests/contrib/test_pandas.py

Original file line number	Diff line number	Diff line change
`@@ -343,7 +343,9 @@ def lstrip_lines(lines: Union[str, List[str]], check=True) -> str:`
`343`	`343`
`344`	`344`
`345`	`345`	`_SKIP_CLOSURE_OBJECTS: Dict[str, Dict[str, Set[str]]] = {`
`346`		`- "globals": {"re": {"_cache"}},`
	`346`	`+ # In onnx, "protobuf" module is imported using "google.protobuf" namespace which results in identifying "google"`
	`347`	`+ # as possible installable requirement which is incorrect. TODO - see if this can be handled in more correct way`
	`348`	`+ "globals": {"re": {"_cache"}, "onnx": {"google"}},`
`347`	`349`	`"nonlocals": {},`
`348`	`350`	`}`
`349`	`351`