Add type annotations to onnxruntime_inference_collection.py (#15364)

### Description

Add type annotations to `onnxruntime_inference_collection.py`



### Motivation and Context

Fixes #15334
This commit is contained in:
Justin Chu 2023-04-05 10:32:49 -07:00 committed by GitHub
parent 764e489a00
commit a96e19abc4
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23

View file

@ -2,16 +2,22 @@
# Copyright (c) Microsoft Corporation. All rights reserved.
# Licensed under the MIT License.
# --------------------------------------------------------------------------
from __future__ import annotations
import collections
import collections.abc
import os
import pathlib
import typing
import warnings
from typing import Any, Sequence
from onnxruntime.capi import _pybind_state as C
if typing.TYPE_CHECKING:
import onnxruntime
def get_ort_device_type(device_type, device_index):
def get_ort_device_type(device_type: str, device_index) -> C.OrtDevice:
if device_type == "cuda":
return C.OrtDevice.cuda()
elif device_type == "cpu":
@ -22,7 +28,11 @@ def get_ort_device_type(device_type, device_index):
raise Exception("Unsupported device type: " + device_type)
def check_and_normalize_provider_args(providers, provider_options, available_provider_names):
def check_and_normalize_provider_args(
providers: Sequence[str, tuple[str, dict[Any, Any]]] | None,
provider_options: Sequence[dict[Any, Any]] | None,
available_provider_names: Sequence[str],
):
"""
Validates the 'providers' and 'provider_options' arguments and returns a
normalized version.
@ -206,8 +216,7 @@ class Session:
# Fallback only once.
self.disable_fallback()
return self._sess.run(output_names, input_feed, run_options)
else:
raise
raise
def run_with_ort_values(self, output_names, input_dict_ort_values, run_options=None):
"""
@ -252,8 +261,7 @@ class Session:
# Fallback only once.
self.disable_fallback()
return invoke(self._sess, output_names, input_dict_ort_values, run_options)
else:
raise
raise
def end_profiling(self):
"""
@ -311,10 +319,17 @@ class InferenceSession(Session):
This is the main class used to run a model.
"""
def __init__(self, path_or_bytes, sess_options=None, providers=None, provider_options=None, **kwargs):
def __init__(
self,
path_or_bytes: str | bytes | os.PathLike,
sess_options: Sequence[onnxruntime.SessionOptions] | None = None,
providers: Sequence[str, tuple[str, dict[Any, Any]]] | None = None,
provider_options: Sequence[dict[Any, Any]] | None = None,
**kwargs,
) -> None:
"""
:param path_or_bytes: filename or serialized ONNX or ORT format model in a byte string
:param sess_options: session options
:param path_or_bytes: Filename or serialized ONNX or ORT format model in a byte string.
:param sess_options: Session options.
:param providers: Optional sequence of providers in order of decreasing
precedence. Values can either be provider names or tuples of
(provider name, options dict). If not provided, then all available
@ -342,13 +357,12 @@ class InferenceSession(Session):
means execute a node using `CUDAExecutionProvider`
if capable, otherwise execute using `CPUExecutionProvider`.
"""
Session.__init__(self)
super().__init__()
if isinstance(path_or_bytes, str):
self._model_path = path_or_bytes
self._model_bytes = None
elif isinstance(path_or_bytes, pathlib.Path):
elif isinstance(path_or_bytes, os.PathLike):
self._model_path = str(path_or_bytes)
self._model_bytes = None
elif isinstance(path_or_bytes, bytes):
@ -374,8 +388,8 @@ class InferenceSession(Session):
self._create_inference_session(self._fallback_providers, None)
# Fallback only once.
self.disable_fallback()
else:
raise
return
raise
def _create_inference_session(self, providers, provider_options, disabled_optimizers=None):
available_providers = C.get_available_providers()
@ -450,7 +464,7 @@ class IOBinding:
This class provides API to bind input/output to a specified device, e.g. GPU.
"""
def __init__(self, session):
def __init__(self, session: Session):
self._iobinding = C.SessionIOBinding(session._sess)
self._numpy_obj_references = {}