mirror of
https://github.com/saymrwulf/onnxruntime.git
synced 2026-07-30 20:18:08 +00:00
Fix InferenceSession API
This commit is contained in:
parent
83c3da3fc0
commit
6d769d47c4
1 changed files with 2 additions and 36 deletions
|
|
@ -76,7 +76,7 @@ class Session:
|
|||
|
||||
def get_session_options(self):
|
||||
"Return the session options. See :class:`onnxruntime.SessionOptions`."
|
||||
return self._sess_options
|
||||
return self._session_options
|
||||
|
||||
def get_inputs(self):
|
||||
"Return the inputs metadata as a list of :class:`onnxruntime.NodeArg`."
|
||||
|
|
@ -193,41 +193,7 @@ class InferenceSession(Session):
|
|||
:param providers: providers to use for session. If empty, will use
|
||||
all available providers.
|
||||
"""
|
||||
self._path_or_bytes = path_or_bytes
|
||||
self._sess_options = sess_options
|
||||
self._load_model(providers)
|
||||
self._enable_fallback = True
|
||||
Session.__init__(self, self._sess)
|
||||
|
||||
def _load_model(self, providers=[]):
|
||||
if isinstance(self._path_or_bytes, str):
|
||||
self._sess = C.InferenceSession(
|
||||
self._sess_options if self._sess_options else C.get_default_session_options(),
|
||||
self._path_or_bytes, True)
|
||||
elif isinstance(self._path_or_bytes, bytes):
|
||||
self._sess = C.InferenceSession(
|
||||
self._sess_options if self._sess_options else C.get_default_session_options(),
|
||||
self._path_or_bytes, False)
|
||||
# elif isinstance(self._path_or_bytes, tuple):
|
||||
# to remove, hidden trick
|
||||
# self._sess.load_model_no_init(self._path_or_bytes[0], providers)
|
||||
else:
|
||||
raise TypeError("Unable to load from type '{0}'".format(type(self._path_or_bytes)))
|
||||
|
||||
self._sess.load_model(providers)
|
||||
|
||||
self._sess_options = self._sess.session_options
|
||||
self._inputs_meta = self._sess.inputs_meta
|
||||
self._outputs_meta = self._sess.outputs_meta
|
||||
self._overridable_initializers = self._sess.overridable_initializers
|
||||
self._model_meta = self._sess.model_meta
|
||||
self._providers = self._sess.get_providers()
|
||||
|
||||
# Tensorrt can fall back to CUDA. All others fall back to CPU.
|
||||
if 'TensorrtExecutionProvider' in C.get_available_providers():
|
||||
self._fallback_providers = ['CUDAExecutionProvider', 'CPUExecutionProvider']
|
||||
else:
|
||||
self._fallback_providers = ['CPUExecutionProvider']
|
||||
Session.__init__(self, path_or_bytes, sess_options, providers)
|
||||
|
||||
|
||||
class IOBinding:
|
||||
|
|
|
|||
Loading…
Reference in a new issue