From a96e19abc4a49943ead82fde2290f219dee833b6 Mon Sep 17 00:00:00 2001 From: Justin Chu Date: Wed, 5 Apr 2023 10:32:49 -0700 Subject: [PATCH] Add type annotations to `onnxruntime_inference_collection.py` (#15364) ### Description Add type annotations to `onnxruntime_inference_collection.py` ### Motivation and Context Fixes #15334 --- .../onnxruntime_inference_collection.py | 46 ++++++++++++------- 1 file changed, 30 insertions(+), 16 deletions(-) diff --git a/onnxruntime/python/onnxruntime_inference_collection.py b/onnxruntime/python/onnxruntime_inference_collection.py index 7a387abb3b..fad1762336 100644 --- a/onnxruntime/python/onnxruntime_inference_collection.py +++ b/onnxruntime/python/onnxruntime_inference_collection.py @@ -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 = {}