Skip to content

Commit

Permalink
Onnx GPU runtime fails to fallback to CPU when GPU is not available/b…
Browse files Browse the repository at this point in the history
…usy (#5304)

* ONNX GPU runtime fails to fallback to CPU when GPU is not available OR busy
#5299

* comments

* Init _fallback_providers before C.InferenceSession

* As per review: Fallback providers order supersedes user's providers order, IF they are included into providers list.

* Code convention fix

* pep8
  • Loading branch information
vladbph authored Oct 3, 2020
1 parent 4721729 commit c20fcf2
Showing 1 changed file with 17 additions and 7 deletions.
24 changes: 17 additions & 7 deletions onnxruntime/python/session.py
Original file line number Diff line number Diff line change
Expand Up @@ -192,9 +192,25 @@ def __init__(self, path_or_bytes, sess_options=None, providers=None, provider_op
self._enable_fallback = True
self._read_config_from_model = os.environ.get('ORT_LOAD_CONFIG_FROM_MODEL') == '1'

self._create_inference_session(providers, provider_options)
try:
self._create_inference_session(providers, provider_options)
except RuntimeError:
if self._enable_fallback:
print("EP Error using {}".format(self._providers))
print("Falling back to {} and retrying.".format(self._fallback_providers))
self._create_inference_session(self._fallback_providers)
# Fallback only once.
self.disable_fallback()
else:
raise

def _create_inference_session(self, providers, provider_options):
# 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_options = self._sess_options if self._sess_options else C.get_default_session_options()
if self._model_path:
sess = C.InferenceSession(session_options, self._model_path, True, self._read_config_from_model)
Expand All @@ -213,12 +229,6 @@ def _create_inference_session(self, providers, provider_options):
self._providers = self._sess.get_providers()
self._provider_options = self._sess.get_provider_options()

# 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']

def _reset_session(self, providers, provider_options):
"release underlying session object."
# meta data references session internal structures
Expand Down

0 comments on commit c20fcf2

Please sign in to comment.