| |
| |
| |
| |
| """ |
| Implements ONNX's backend API. |
| """ |
|
|
| from onnx.backend.base import BackendRep |
|
|
| from onnxruntime import RunOptions |
|
|
| |
| |
| |
| _ALLOWED_RUN_OPTIONS = frozenset( |
| { |
| "log_severity_level", |
| "log_verbosity_level", |
| "logid", |
| "only_execute_path_to_fetches", |
| } |
| ) |
|
|
|
|
| class OnnxRuntimeBackendRep(BackendRep): |
| """ |
| Wraps an :class:`onnxruntime.InferenceSession` to implement ONNX's |
| :class:`onnx.backend.base.BackendRep` interface for running predictions. |
| """ |
|
|
| def __init__(self, session): |
| """ |
| :param session: :class:`onnxruntime.InferenceSession` |
| """ |
| self._session = session |
|
|
| def run(self, inputs, **kwargs): |
| """ |
| Computes the prediction. |
| See :meth:`onnxruntime.InferenceSession.run`. |
| |
| :param inputs: a list of input arrays (one per model input) or a single |
| array when the model has exactly one input |
| :param kwargs: only a safe subset of :class:`onnxruntime.RunOptions` attributes are |
| accepted; see ``_ALLOWED_RUN_OPTIONS`` for the list |
| :return: list of output arrays |
| """ |
|
|
| options = RunOptions() |
| for k, v in kwargs.items(): |
| if k in _ALLOWED_RUN_OPTIONS: |
| setattr(options, k, v) |
| elif hasattr(options, k): |
| raise RuntimeError( |
| f"RunOptions attribute '{k}' is not permitted via the backend API. " |
| f"Allowed attributes: {', '.join(sorted(_ALLOWED_RUN_OPTIONS))}" |
| ) |
| |
|
|
| if isinstance(inputs, list): |
| inps = {} |
| for i, inp in enumerate(self._session.get_inputs()): |
| inps[inp.name] = inputs[i] |
| outs = self._session.run(None, inps, options) |
| if isinstance(outs, list): |
| return outs |
| else: |
| output_names = [o.name for o in self._session.get_outputs()] |
| return [outs[name] for name in output_names] |
| else: |
| inp = self._session.get_inputs() |
| if len(inp) != 1: |
| raise RuntimeError(f"Model expect {len(inp)} inputs") |
| inps = {inp[0].name: inputs} |
| return self._session.run(None, inps, options) |
|
|
