saymrwulf/onnxruntime
1
0
Fork 0
mirror of https://github.com/saymrwulf/onnxruntime.git synced 2026-07-07 17:15:29 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

fix windows cuda support for python 3.8 + (#6046)

Browse source 
* fix

* noqa

* fix.

* remove unused import
This commit is contained in:
George Wu 2020-12-07 10:09:22 -08:00 committed by GitHub
parent 7cebf76a46
commit 020efc9002
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
2 changed files with 11 additions and 14 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

11
onnxruntime/__init__.py
View file

@ -10,6 +10,17 @@ or the `Github project <https://github.com/microsoft/onnxruntime/>`_.
__version__ = "1.5.2"
__author__ = "Microsoft"
import os
import platform
import sys
# Python 3.8 (and later) on Windows doesn't search system PATH when loading DLLs,
# so CUDA location needs to be specified explicitly. This needs to be done before importing
# onnxruntime.capi._pybind_state
if "CUDA_PATH" in os.environ and platform.system() == "Windows" and sys.version_info >= (3, 8):
cuda_bin_dir = os.path.join(os.environ["CUDA_PATH"], "bin")
os.add_dll_directory(cuda_bin_dir)
from onnxruntime.capi._pybind_state import get_all_providers, get_available_providers, get_device, set_seed, \
RunOptions, SessionOptions, set_default_logger_severity, enable_telemetry_events, disable_telemetry_events, \
NodeArg, ModelMetadata, GraphOptimizationLevel, ExecutionMode, ExecutionOrder, OrtDevice, SessionIOBinding, \

14
onnxruntime/python/_pybind_state.py
View file

@ -5,23 +5,9 @@
import os
import platform
import sys
import warnings
import onnxruntime.capi._ld_preload # noqa: F401
# Python 3.8 (and later) on Windows doesn't search system PATH when loading DLLs,
# so CUDA location needs to be specified explicitly.
if platform.system() == "Windows" and sys.version_info >= (3, 8):
CUDA_VERSION = "10.2"
CUDNN_VERSION = "8"
cuda_env_variable = "CUDA_PATH_V" + CUDA_VERSION.replace(".", "_")
if cuda_env_variable not in os.environ:
raise ImportError("CUDA Toolkit %s not installed on the machine." % CUDA_VERSION)
cuda_bin_dir = os.path.join(os.environ[cuda_env_variable], "bin")
if not os.path.isfile(os.path.join(cuda_bin_dir, "cudnn64_%s.dll" % CUDNN_VERSION)):
raise ImportError("cuDNN %s not installed on the machine." % CUDNN_VERSION)
os.add_dll_directory(cuda_bin_dir)
try:
from onnxruntime.capi.onnxruntime_pybind11_state import * # noqa
except ImportError as e: