onnx-diagnostic 0.8.0__py3-none-any.whl

onnx_diagnostic/reference/ort_evaluator.py

@@ -0,0 +1,652 @@
+from typing import Any, Dict, List, Optional, Sequence, Set, Tuple, Union
+import numpy as np
+from onnx import (
+    AttributeProto,
+    GraphProto,
+    FunctionProto,
+    ModelProto,
+    NodeProto,
+    TypeProto,
+    ValueInfoProto,
+    helper as oh,
+    load,
+    save as onnx_save,
+    shape_inference as shi,
+)
+from onnx.defs import onnx_opset_version
+import onnxruntime
+from ..helpers import string_type
+from ..helpers.onnx_helper import pretty_onnx, dtype_to_tensor_dtype, to_array_extended
+from ..helpers.ort_session import (
+    InferenceSessionForTorch,
+    InferenceSessionForNumpy,
+    _InferenceSession,
+)
+from ..helpers.torch_helper import to_tensor
+from .report_results_comparison import ReportResultComparison
+from .evaluator import ExtendedReferenceEvaluator
+
+
+PROTO = (FunctionProto, ModelProto, GraphProto, NodeProto)
+Proto = Union[FunctionProto, ModelProto, GraphProto, NodeProto]
+
+
+class OnnxruntimeEvaluator:
+    """
+    This class loads an onnx model and the executes one by one the nodes
+    with onnxruntime. This class is mostly meant for debugging.
+
+    :param proto: proto or filename
+    :param session_options: options
+    :param providers: providers
+    :param nvtx: enable nvidia events
+    :param providers: `None`, `"CPU"`, `"CUDA"` or a list of providers
+    :param graph_optimization_level: see :class:`onnxruntime.SessionOptions`
+    :param log_severity_level: see :class:`onnxruntime.SessionOptions`
+    :param log_verbosity_level: see :class:`onnxruntime.SessionOptions`
+    :param optimized_model_filepath:  see :class:`onnxruntime.SessionOptions`
+    :param disable_aot_function_inlining:  see :class:`onnxruntime.SessionOptions`
+    :param use_training_api: use onnxruntime-traning API
+    :param verbose: verbosity
+    :param local_functions: additional local function
+    :param ir_version: ir version to use when unknown
+    :param opsets: opsets to use when unknown
+    :param whole: if True, do not split node by node
+    :param torch_or_numpy: force the use of one of them, True for torch,
+        False for numpy, None to let the class choose
+    """
+
+    def __init__(
+        self,
+        proto: Union[str, Proto, "OnnxruntimeEvaluator"],
+        session_options: Optional[onnxruntime.SessionOptions] = None,
+        providers: Optional[Union[str, List[str]]] = None,
+        nvtx: bool = False,
+        enable_profiling: bool = False,
+        graph_optimization_level: Union[onnxruntime.GraphOptimizationLevel, bool] = None,
+        log_severity_level: Optional[int] = None,
+        log_verbosity_level: Optional[int] = None,
+        optimized_model_filepath: Optional[str] = None,
+        disable_aot_function_inlining: Optional[bool] = None,
+        use_training_api: bool = False,
+        verbose: int = 0,
+        local_functions: Optional[
+            Dict[Tuple[str, str], Union[Proto, "OnnxruntimeEvaluator"]]
+        ] = None,
+        ir_version: int = 10,
+        opsets: Optional[Union[int, Dict[str, int]]] = None,
+        whole: bool = False,
+        torch_or_numpy: Optional[bool] = None,
+    ):
+        if isinstance(proto, str):
+            self.proto: Proto = load(proto)
+        elif isinstance(proto, OnnxruntimeEvaluator):
+            assert isinstance(
+                proto.proto, PROTO
+            ), f"Unexpected type for proto.proto {type(proto.proto)}"
+            self.proto = proto.proto
+        else:
+            self.proto = proto
+        assert isinstance(
+            self.proto, PROTO
+        ), f"Unexpected type for self.proto {type(self.proto)}"
+
+        self._cache: Dict[
+            Any, Tuple[Proto, Union["OnnxruntimeEvaluator", _InferenceSession]]  # noqa: UP037
+        ] = {}
+        self.ir_version = ir_version
+        self.opsets = opsets
+        self.session_kwargs: Dict[str, Any] = dict(
+            session_options=session_options,
+            providers=providers,
+            nvtx=nvtx,
+            enable_profiling=enable_profiling,
+            graph_optimization_level=graph_optimization_level,
+            log_severity_level=log_severity_level,
+            log_verbosity_level=log_verbosity_level,
+            optimized_model_filepath=optimized_model_filepath,
+            disable_aot_function_inlining=disable_aot_function_inlining,
+            use_training_api=use_training_api,
+        )
+        self.to_tensor_or_array = to_array_extended if not torch_or_numpy else to_tensor
+
+        self.verbose = verbose
+        self.torch_or_numpy = torch_or_numpy
+        self.sess_: Optional[_InferenceSession] = None
+        if whole:
+            self.nodes: Optional[List[NodeProto]] = None
+            self.rt_inits_: Optional[Dict[str, Any]] = None
+            self.rt_nodes_: Optional[List[NodeProto]] = None
+        else:
+            self.nodes = (
+                [self.proto]
+                if isinstance(self.proto, NodeProto)
+                else (
+                    list(
+                        self.proto.graph.node
+                        if hasattr(self.proto, "graph")
+                        else self.proto.node
+                    )
+                )
+            )
+            self.rt_inits_ = (
+                {
+                    init.name: self.to_tensor_or_array(init)
+                    for init in self.proto.graph.initializer
+                }
+                if hasattr(self.proto, "graph")
+                else {}
+            )
+            self.rt_nodes_ = self.nodes.copy()
+
+        self.local_functions: Dict[Tuple[str, str], "OnnxruntimeEvaluator"] = (  # noqa: UP037
+            {(f.domain, f.name): self.__class__(f) for f in self.proto.functions}
+            if hasattr(self.proto, "functions")
+            else {}
+        )
+        if local_functions:
+            self.local_functions.update(local_functions)
+        self.garbage_collector = self._build_garbage_collector() if self.rt_nodes_ else {}
+
+    @property
+    def input_names(self) -> List[str]:
+        "Returns input names."
+        assert self.proto, "self.proto is empty"
+        if isinstance(self.proto, NodeProto):
+            assert isinstance(
+                self.nodes, list
+            ), f"Unexpected type {type(self.nodes)} for self.nodes"
+            return self.nodes[0].input
+        return [
+            getattr(o, "name", o)
+            for o in (
+                self.proto.graph.input if hasattr(self.proto, "graph") else self.proto.input
+            )
+        ]
+
+    @property
+    def output_names(self) -> List[str]:
+        "Returns output names."
+        assert self.proto, "self.proto is empty"
+        if isinstance(self.proto, NodeProto):
+            assert isinstance(
+                self.nodes, list
+            ), f"Unexpected type {type(self.nodes)} for self.nodes"
+            return self.nodes[0].output
+        return [
+            getattr(o, "name", o)
+            for o in (
+                self.proto.graph.output if hasattr(self.proto, "graph") else self.proto.output
+            )
+        ]
+
+    @property
+    def input_types(self) -> List[TypeProto]:
+        "Returns input types."
+        if not isinstance(self.proto, (ModelProto, GraphProto)):
+            raise ValueError(f"Cannot guess input types for type {type(self.proto)}")
+        g = self.proto.graph if hasattr(self.proto, "graph") else self.proto
+        return [i.type for i in g.input]
+
+    @property
+    def output_types(self) -> List[TypeProto]:
+        "Returns output types."
+        if not isinstance(self.proto, (ModelProto, GraphProto)):
+            raise ValueError(f"Cannot guess output types for type {type(self.proto)}")
+        g = self.proto.graph if hasattr(self.proto, "graph") else self.proto
+        return [i.type for i in g.output]
+
+    def _log_arg(self, a: Any) -> Any:
+        if isinstance(a, (str, int, float)):
+            return a
+        device = f"D{a.get_device()}:" if hasattr(a, "detach") else ""
+        if hasattr(a, "shape"):
+            prefix = "A:" if hasattr(a, "astype") else "T:"
+            if self.verbose < 4:  # noqa: PLR2004
+                return f"{prefix}{device}{a.dtype}:{a.shape} in [{a.min()}, {a.max()}]"
+            elements = a.ravel().tolist()
+            if len(elements) > 10:  # noqa: PLR2004
+                elements = elements[:10]
+                return f"{prefix}{device}{a.dtype}:{a.shape}:{','.join(map(str, elements))}..."
+            return f"{prefix}{device}{a.dtype}:{a.shape}:{elements}"
+        if hasattr(a, "append"):
+            return ", ".join(map(self._log_arg, a))
+        return a
+
+    def _log(self, level: int, pattern: str, *args: Any) -> None:
+        if level < self.verbose:
+            new_args = [self._log_arg(a) for a in args]
+            print(pattern % tuple(new_args))
+
+    def _is_local_function(self, node: NodeProto) -> bool:
+        return (node.domain, node.op_type) in self.local_functions
+
+    def run(
+        self,
+        outputs: Optional[List[str]],
+        feed_inputs: Dict[str, Any],
+        intermediate: bool = False,
+        report_cmp: Optional[ReportResultComparison] = None,
+    ) -> Union[Dict[str, Any], List[Any]]:
+        """
+        Runs the model.
+        It only works with numpy arrays.
+
+        :param outputs: required outputs or None for all
+        :param feed_inputs: inputs
+        :param intermediate: returns all output instead of the last ones
+        :param report_cmp: used as a reference,
+            every intermediate results is compare to every existing one,
+            if not empty, it is an instance of
+            :class:`onnx_diagnostic.reference.ReportResultComparison`
+        :return: outputs, as a list if return_all is False,
+            as a dictionary if return_all is True
+        """
+        if self.rt_nodes_ is None:
+            # runs a whole
+            if self.sess_ is None:
+                assert self.proto, "self.proto is empty"
+                _, self.sess_ = self._get_sess(self.proto, list(feed_inputs.values()))
+            assert self.sess_, "mypy not happy"
+            return self.sess_.run(outputs, feed_inputs)
+        if outputs is None:
+            outputs = self.output_names
+        results: Dict[str, Any] = (self.rt_inits_ or {}).copy()
+
+        for k, v in results.items():
+            self._log(2, " +C %s: %s", k, v)
+        for k, v in feed_inputs.items():
+            assert not isinstance(v, str), f"Unexpected type str for {k!r}"
+            self._log(2, " +I %s: %s", k, v)
+            results[k] = v
+
+        for i_node, node in enumerate(self.rt_nodes_ or []):
+            self._log(1, "%s(%s) -> %s", node.op_type, node.input, node.output)
+            for i in node.input:
+                if i != "" and i not in results:
+                    raise RuntimeError(
+                        f"Unable to find input {i!r} in known results {sorted(results)}, "
+                        f"self.rt_inits_ has {sorted((self.rt_inits_ or {}))}, "
+                        f"feed_inputs has {sorted(feed_inputs)}."
+                    )
+            inputs = [(results[i] if i != "" else None) for i in node.input]
+            if node.op_type == "If" and node.domain == "":
+                outputs = self._run_if(node, inputs, results)
+            elif node.op_type in {"Scan", "Loop"} and node.domain == "":
+                outputs = self._run_scan(node, inputs, results)
+            elif self._is_local_function(node):
+                outputs = self._run_local(node, inputs, results)
+            else:
+                outputs = self._run(node, inputs, results)
+            for name, value in zip(node.output, outputs):
+                if name == "":
+                    continue
+                self._log(2, " + %s: %s", name, value)  # type: ignore[arg-type]
+                assert isinstance(name, str), f"unexpected type for name {type(name)}"
+                results[name] = value
+            if report_cmp:
+                reported = report_cmp.report(dict(zip(node.output, outputs)))
+                if self.verbose > 1:
+                    print(f"  -- report {len(reported)} comparisons")
+            if not intermediate:
+                self._clean_unused_inplace(i_node, node, results)
+
+        if intermediate:
+            return results
+        output_names = self.output_names
+        for name in output_names:
+            if name == "":
+                continue
+            if name not in results:
+                raise RuntimeError(
+                    f"Unable to find output name {name!r} "
+                    f"in {sorted(results)}, proto is\n{pretty_onnx(self.proto)}"
+                )
+        return [results[name] for name in output_names if name != ""]
+
+    def _build_garbage_collector(self) -> Dict[str, int]:
+        """
+        Memorizes the results not needed anymore for every node.
+        Returns a dictionary with the last node using the results.
+        """
+        needed = {}
+        for i, node in enumerate(self.rt_nodes_ or []):
+            for name in node.input:
+                needed[name] = i
+            if node.op_type in {"Scan", "If", "Loop"}:
+                hidden = self._get_hidden_node_inputs(node)
+                for name in hidden:
+                    needed[name] = i
+        if isinstance(self.proto, ModelProto):
+            for o in self.proto.graph.output:
+                needed[o.name] = len(self.rt_nodes_ or [])
+        elif isinstance(self.proto, GraphProto):
+            for o in self.proto.output:
+                needed[o.name] = len(self.rt_nodes_ or [])
+        elif isinstance(self.proto, FunctionProto):
+            for o in self.proto.output:
+                needed[o] = len(self.rt_nodes_ or [])
+        return needed
+
+    def _clean_unused_inplace(self, i_node: int, node: NodeProto, results: Dict[str, Any]):
+        """
+        Cleans all results not needed anymore. Some models requires to clean the memory
+        to be able to run.
+        """
+        if not self.garbage_collector:
+            return
+        for name in node.input:
+            if self.garbage_collector[name] == i_node and name in results:
+                if self.verbose:
+                    t = results[name]
+                    print(f" - deletes: {name} - {t.dtype}:{t.shape}")
+                del results[name]
+        if node.op_type in {"Scan", "If", "Loop"}:
+            hidden = self._get_hidden_node_inputs(node)
+            for name in hidden:
+                if self.garbage_collector[name] == i_node and name in results:
+                    if self.verbose:
+                        t = results[name]
+                        print(f" - deletes: {name} - {t.dtype}:{t.shape}")
+                    del results[name]
+
+    def _make_model_proto(
+        self,
+        nodes: Sequence[NodeProto],
+        vinputs: Sequence[ValueInfoProto],
+        voutputs: Sequence[ValueInfoProto],
+    ) -> ModelProto:
+        onx = oh.make_model(
+            oh.make_graph(nodes, "-", vinputs, voutputs),
+            ir_version=getattr(self.proto, "ir_version", self.ir_version),
+            functions=getattr(self.proto, "functions", None),
+        )
+        del onx.opset_import[:]
+        if hasattr(self.proto, "opset_import"):
+            onx.opset_import.extend(self.proto.opset_import)
+        elif self.opsets:
+            if isinstance(self.opsets, int):
+                onx.opset_import.append(oh.make_opsetid("", self.opsets))
+            else:
+                onx.opset_import.extend(
+                    [oh.make_opsetid(k, v) for k, v in self.opsets.items()]
+                )
+        else:
+            onx.opset_import.append(oh.make_opsetid("", onnx_opset_version()))
+
+        # That helps fixing bugs.
+        onx = shi.infer_shapes(onx)
+        return onx
+
+    @classmethod
+    def _get_hidden_inputs(self, graph: GraphProto) -> Set[str]:
+        """
+        Returns the hidden inputs (inputs coming from an upper context)
+        used by a subgraph.
+        """
+        hidden = set()
+        memo = set(i.name for i in graph.initializer)
+        memo |= set(i.name for i in graph.sparse_initializer)
+        for node in graph.node:
+            for i in node.input:
+                if i not in memo:
+                    hidden.add(i)
+            for att in node.attribute:
+                if att.type == AttributeProto.GRAPH and att.g:
+                    hid = self._get_hidden_inputs(att.g)
+                    less = set(h for h in hid if h not in memo)
+                    hidden |= less
+            memo |= set(node.output)
+        return hidden
+
+    @classmethod
+    def _get_hidden_node_inputs(self, node: NodeProto) -> Set[str]:
+        """Calls multiple _get_hidden_inputs on every attribute."""
+        if node.op_type not in {"Loop", "Scan", "If"}:
+            return set()
+        hidden = set()
+        for att in node.attribute:
+            if att.type == AttributeProto.GRAPH:
+                hidden |= self._get_hidden_inputs(att.g)
+        return hidden - (hidden & set(node.input))
+
+    def _get_sess(
+        self, node: Union[ModelProto, NodeProto], inputs: List[Any]
+    ) -> Tuple[ModelProto, _InferenceSession]:
+        if isinstance(node, ModelProto):
+            onx = node
+        else:
+            assert isinstance(node, NodeProto), f"Unexpected type {type(node)} for node"
+            if node.op_type == "Constant":
+                # We force the type to be a boolean.
+                ref = ExtendedReferenceEvaluator(node)
+                cst = ref.run(None, {})[0]
+                vinputs: List[ValueInfoProto] = []
+                voutputs = [
+                    oh.make_tensor_value_info(
+                        node.output[0], dtype_to_tensor_dtype(cst.dtype), cst.shape
+                    )
+                ]
+            else:
+                unique_names = set()
+                vinputs = []
+                for i, it in zip(node.input, inputs):
+                    if i == "" or i in unique_names:
+                        continue
+                    unique_names.add(i)
+                    value = oh.make_tensor_value_info(
+                        i, dtype_to_tensor_dtype(it.dtype), it.shape
+                    )
+                    vinputs.append(value)
+
+                # no need to run shape inference
+                voutputs = [oh.make_value_info(o, TypeProto()) for o in node.output]
+
+            onx = self._make_model_proto([node], vinputs, voutputs)
+
+        cls = (
+            InferenceSessionForNumpy
+            if any(isinstance(i, np.ndarray) for i in inputs)
+            and (not isinstance(self.torch_or_numpy, bool) or not self.torch_or_numpy)
+            else InferenceSessionForTorch
+        )
+        try:
+            sess = cls(onx, **self.session_kwargs)
+        except (
+            onnxruntime.capi.onnxruntime_pybind11_state.Fail,
+            onnxruntime.capi.onnxruntime_pybind11_state.InvalidGraph,
+            onnxruntime.capi.onnxruntime_pybind11_state.InvalidArgument,
+        ) as e:
+            onnx_save(onx, "_debug_OnnxruntimeEvaluator_last_failure.onnx")
+            raise RuntimeError(
+                f"Unable to infer a session with inputs\n{string_type(inputs)}"
+                f"\ndue to {e}\n{pretty_onnx(onx)}"
+            ) from e
+        return onx, sess
+
+    def _get_sess_init_subgraph(
+        self, node: NodeProto, inputs: List[Any], context: Dict[str, Any], g: GraphProto
+    ) -> List[Any]:
+        unique_names = set()
+        vinputs = []
+        for i, it in zip(node.input, inputs):
+            if i == "" or i in unique_names:
+                continue
+            unique_names.add(i)
+            value = oh.make_tensor_value_info(i, dtype_to_tensor_dtype(it.dtype), it.shape)
+            vinputs.append(value)
+
+        reduced_set = self._get_hidden_inputs(g)
+        for i, v in context.items():
+            if i in reduced_set and i not in unique_names:
+                unique_names.add(i)
+                value = oh.make_tensor_value_info(i, dtype_to_tensor_dtype(v.dtype), v.shape)
+                vinputs.append(value)
+        return vinputs
+
+    def _get_sess_if(
+        self, node: NodeProto, branch: str, inputs: List[Any], context: Dict[str, Any]
+    ) -> Tuple[ModelProto, "OnnxruntimeEvaluator"]:
+        g = None
+        for att in node.attribute:
+            if att.name == branch:
+                g = att.g
+        assert g, f"Missing attribute {branch!r}"
+        vinputs = self._get_sess_init_subgraph(node, inputs, context, g)
+
+        voutputs = g.output
+
+        identities = [
+            oh.make_node("Identity", [iname], [ginput.name])
+            for iname, ginput in zip(node.input, g.input)
+        ]
+
+        onx = self._make_model_proto([*identities, *g.node], vinputs, voutputs)
+        sess = OnnxruntimeEvaluator(
+            onx,
+            local_functions=self.local_functions,
+            verbose=self.verbose,
+            ir_version=self.ir_version,
+            opsets=self.opsets,
+            torch_or_numpy=self.torch_or_numpy,
+            **self.session_kwargs,
+        )
+        return onx, sess
+
+    def _get_sess_local(
+        self, node: NodeProto, inputs: List[Any]
+    ) -> Tuple[FunctionProto, "OnnxruntimeEvaluator"]:
+        ev = self.local_functions[node.domain, node.op_type]
+        sess = OnnxruntimeEvaluator(
+            ev,
+            local_functions=self.local_functions,
+            verbose=self.verbose,
+            ir_version=self.ir_version,
+            opsets=self.opsets,
+            torch_or_numpy=self.torch_or_numpy,
+            **self.session_kwargs,
+        )
+        return ev.proto, sess
+
+    def _run(self, node: NodeProto, inputs: List[Any], results: Dict[str, Any]) -> List[Any]:
+        """Runs a node."""
+        types = [(None if a is None else (a.dtype, a.shape)) for a in inputs]
+        key = (id(node), *types)
+        if key in self._cache:
+            sess = self._cache[key][1]
+        else:
+            onx, sess = self._get_sess(node, inputs)
+            self._cache[key] = onx, sess
+
+        feeds = dict(zip(node.input, inputs))
+        if "" in feeds:
+            feeds[""] = np.array([0], dtype=np.float32)
+
+        assert hasattr(sess, "run"), f"Missing method run for type {type(sess)}"
+        outputs = list(sess.run(None, feeds))
+        assert isinstance(outputs, list), f"Unexpected type for outputs {type(outputs)}"
+        return outputs
+
+    def _run_if(
+        self, node: NodeProto, inputs: List[Any], results: Dict[str, Any]
+    ) -> List[Any]:
+        """Runs a node If."""
+        feeds = dict(zip(node.input, inputs))
+        feeds.update(results)
+        if feeds[node.input[0]]:
+            name = "then_branch"
+        else:
+            name = "else_branch"
+
+        key = (id(node), name)
+        if key in self._cache:
+            sess = self._cache[key][1]
+        else:
+            self._cache[key] = _onx, sess = self._get_sess_if(node, name, inputs, results)
+
+        assert hasattr(sess, "run"), f"Missing method run for type {type(sess)}"
+        feeds = {name: results[name] for name in sess.input_names}
+        outputs = sess.run(None, feeds)
+        assert isinstance(outputs, list), f"Unexpected type for outputs {type(outputs)}"
+        return outputs
+
+    def _get_sess_scan(
+        self, node: NodeProto, branch: str, inputs: List[Any], context: Dict[str, Any]
+    ) -> Tuple[ModelProto, "OnnxruntimeEvaluator"]:
+        g = None
+        for att in node.attribute:
+            if att.name == branch:
+                g = att.g
+        assert g, f"Missing attribute {branch!r}"
+        vinputs = self._get_sess_init_subgraph(node, inputs, context, g)
+
+        begin = 0 if node.op_type == "Scan" else 1
+        voutputs = []
+        for name, _goutput in zip(node.output, g.output[begin:]):
+            v = ValueInfoProto()
+            # v.ParseFromString(goutput.SerializeToString())
+            v.name = name
+            voutputs.append(v)
+
+        # identities = []
+        # for iname, ginput in zip(node.input, g.input):
+        #    identities.append(oh.make_node("Identity", [iname], [ginput.name]))
+
+        onx = self._make_model_proto([node], vinputs, voutputs)
+        sess = OnnxruntimeEvaluator(
+            onx,
+            local_functions=self.local_functions,
+            verbose=self.verbose,
+            ir_version=self.ir_version,
+            opsets=self.opsets,
+            torch_or_numpy=self.torch_or_numpy,
+            whole=True,
+            **self.session_kwargs,
+        )
+        return onx, sess
+
+    def _run_scan(
+        self, node: NodeProto, inputs: List[Any], results: Dict[str, Any]
+    ) -> List[Any]:
+        """Runs a node Scan."""
+        feeds = dict(zip(node.input, inputs))
+        feeds.update(results)
+        name = "body"
+        key = (id(node), name)
+        if key in self._cache:
+            sess = self._cache[key][1]
+        else:
+            self._cache[key] = _onx, sess = self._get_sess_scan(node, name, inputs, results)
+
+        assert hasattr(sess, "run"), f"Missing method run for type {type(sess)}"
+        feeds = {name: results[name] for name in sess.input_names}
+        outputs = sess.run(None, feeds)
+        assert isinstance(outputs, list), f"Unexpected type for outputs {type(outputs)}"
+        return outputs
+
+    def _run_local(
+        self, node: NodeProto, inputs: List[Any], results: Dict[str, Any]
+    ) -> List[Any]:
+        """Runs a node."""
+        types = [(None if a is None else (a.dtype, a.shape)) for a in inputs]
+        key = (id(node), *types)
+        if key in self._cache:
+            sess = self._cache[key][1]
+        else:
+            onx, sess = self._get_sess_local(node, inputs)
+            self._cache[key] = onx, sess
+
+        replace = dict(zip(node.input, sess.input_names))
+        assert len(node.input) == len(sess.input_names), (
+            f"Input mismatch: input_names={sess.input_names}, "
+            f"replace={replace}, "
+            f"type(self.proto)={type(self.proto)}, and node=\n{node}"
+        )
+        feeds = {replace[i]: v for i, v in zip(node.input, inputs)}
+        if "" in feeds:
+            feeds[""] = np.array([0], dtype=np.float32)
+
+        assert hasattr(sess, "run"), f"Missing method run for type {type(sess)}"
+        outputs = sess.run(None, feeds)
+        assert isinstance(outputs, list), f"Unexpected type for outputs {type(outputs)}"
+        return outputs

onnx_diagnostic/reference/quantized_tensor.py

@@ -0,0 +1,46 @@
+import numpy as np
+
+
+class QuantizedTensor:
+    """
+    Quantizes a vector in range [0, 255].
+
+    :param tensor: original tensor
+    """
+
+    def __init__(self, tensor):
+        _min = tensor.min()
+        _max = tensor.max()
+        _min = min(_min, 0)
+        _max = max(_max, 0)
+        qmin = 0
+        qmax = 255
+
+        self.scale_ = np.array((_max - _min) / (qmax - qmin), dtype=tensor.dtype)
+        initial_zero_point = qmin - _min / self.scale_
+        self.zero_point_ = np.array(
+            int(max(qmin, min(qmax, initial_zero_point))), dtype=np.uint8
+        )
+        self.quantized_ = np.maximum(
+            0, np.minimum(qmax, (tensor / self.scale_).astype(int) + self.zero_point_)
+        ).astype(self.zero_point_.dtype)
+
+    @property
+    def shape(self):
+        "accessor"
+        return self.quantized_.shape
+
+    @property
+    def scale(self):
+        "accessor"
+        return self.scale_
+
+    @property
+    def zero_point(self):
+        "accessor"
+        return self.zero_point_
+
+    @property
+    def qtensor(self):
+        "accessor"
+        return self.quantized_